import os from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler import torch import numpy as np from pipeline_stable_diffusion_controlnet_inpaint import * from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import * import random # model2 controlnet = ControlNetModel.from_pretrained("./models/control_any5_openpose", torch_dtype=torch.float16) pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained("./models/Any-inpainting", controlnet=controlnet, torch_dtype=torch.float16) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) def make_inpaint_condition(image, image_mask): image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 image_mask = np.array(image_mask.convert("L")) assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size" image[image_mask > 128] = -1.0 # set as masked pixel image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) image = torch.from_numpy(image) return image def generate_image(prompt:str, negative_prompt:str, openpose_image, original_image, mask_image): # model1 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16, cache_dir='./models') pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( "./models/Any-inpainting", controlnet=controlnet, torch_dtype=torch.float16, cache_dir='./models' ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) pipe.to('cuda') a = random.randint(10000,90000) generator = torch.manual_seed(a) # control_image = make_inpaint_condition(original_image, mask_image) # images = [openpose_image, control_image] image = pipe( prompt=prompt, # images, image=original_image, control_image=openpose_image, mask_image=mask_image, num_inference_steps=20, generator=generator, negative_prompt=negative_prompt, # controlnet_conditioning_scale=[1.0, 0.8], ).images[0] return image def generate_image_sketch(prompt: str, negative_prompt: str, openpose_image, original_image, mask_image): b = random.randint(10000, 90000) generator = torch.manual_seed(b) # model2 controlnet1 = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_scribble", torch_dtype=torch.float16, cache_dir='./models') pipe1 = StableDiffusionControlNetInpaintPipeline.from_pretrained( "./models/Any-inpainting", controlnet=controlnet1, torch_dtype=torch.float16, cache_dir='./models' ) pipe1.scheduler = UniPCMultistepScheduler.from_config(pipe1.scheduler.config) pipe1.to('cuda') image = pipe1( prompt=prompt, # images, image=original_image, control_image=openpose_image, mask_image=mask_image, num_inference_steps=20, generator=generator, negative_prompt=negative_prompt, # controlnet_conditioning_scale=[1.0, 0.8], ).images[0] return [image]