Update generate_img.py

generate_img.py

@@ -5,24 +5,6 @@ import numpy as np
 from pipeline_stable_diffusion_controlnet_inpaint import *
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import *
 import random
-#model1
-controlnet =  ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16)
-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')
-
-# model2
-controlnet1 =  ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_scribble",torch_dtype=torch.float16)
-
-pipe1 = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-    "./models/Any-inpainting", controlnet=controlnet1, torch_dtype=torch.float16
-)
-
-pipe1.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-pipe1.to('cuda')
 
 def make_inpaint_condition(image, image_mask):
     image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
@@ -35,7 +17,16 @@ def make_inpaint_condition(image, image_mask):
 
 
 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)
@@ -59,8 +50,16 @@ def generate_image(prompt:str, negative_prompt:str, openpose_image, original_ima
 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)
-    # control_image = make_inpaint_condition(original_image, mask_image)
-    # images = [openpose_image, control_image]
+    # 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,