xinsir6 / ControlNetPlus

ControlNet++: All-in-one ControlNet for image generations and editing!
Apache License 2.0
1.64k stars 33 forks source link

inpainting doesnt works #42

Closed AbhinavJangra29 closed 1 month ago

AbhinavJangra29 commented 1 month ago

this is the code i am using:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import sys
sys.path.append('..')
import cv2
import copy
import torch
import random
import numpy as np
from PIL import Image
from mask import get_mask_generator
from diffusers.utils import load_image
from diffusers import EulerAncestralDiscreteScheduler, AutoencoderKL
from models.controlnet_union import ControlNetModel_Union
from pipeline.pipeline_controlnet_union_inpaint_sd_xl import StableDiffusionXLControlNetUnionInpaintPipeline

device=torch.device('cuda:0')

eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

controlnet_model = ControlNetModel_Union.from_pretrained(r"controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16, use_safetensors=True)

pipe = StableDiffusionXLControlNetUnionInpaintPipeline.from_pretrained(
    "SG161222/RealVisXL_V4.0", controlnet=controlnet_model, 
    vae=vae,
    torch_dtype=torch.float16,
    #scheduler=ddim_scheduler,
    scheduler=eulera_scheduler,
)

pipe = pipe.to(device)

import numpy as np
import cv2
from PIL import Image
import random
import copy

def HWC3(x):
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y

# Prompt and negative prompt
prompt = "dog sitting"
negative_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'

# Seed for random number generator
seed = random.randint(0, 2147483647)

# The original image you want to repaint
original_img = cv2.imread("img.jpg")

# The mask image (single channel)
mask = cv2.imread("mask3.jpg", cv2.IMREAD_GRAYSCALE)

# Resizing the original image and mask
height, width, _ = original_img.shape
ratio = np.sqrt(1024. * 1024. / (width * height))
W, H = int(width * ratio) // 8 * 8, int(height * ratio) // 8 * 8
original_img = cv2.resize(original_img, (W, H))
original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)

# Resize the mask to match the original image
mask = cv2.resize(mask, (W, H))

# Prepare controlnet_img
controlnet_img = copy.deepcopy(original_img)
controlnet_img[mask > 0] = 0

# Convert mask to three channels
mask = HWC3(mask)

# Convert images to PIL format
controlnet_img = Image.fromarray(controlnet_img)
original_img = Image.fromarray(original_img)
mask = Image.fromarray(mask)

# Setting width and height
width, height = W, H

images = pipe(prompt=prompt,
            image=original_img,
            mask_image=mask,
            control_image_list=[0, 0, 0, 0, 0, 0, 0, controlnet_img], 
            negative_prompt=negative_prompt,
            # generator=generator,
            width=width, 
            height=height,
            num_inference_steps=30,
            union_control=True,
            union_control_type=torch.Tensor([0, 0, 0, 0, 0, 0, 0, 1]),
            ).images

controlnet_img.save("control_inpainting.jpg")
images[0].save(f"masti.png")

input image: image

mask used: image

controlnet image: image

final inpaint image: image

zakajd commented 1 day ago

@AbhinavJangra29 did you figure out what was wrong?