Update generate_img.py

generate_img.py

@@ -7,10 +7,16 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet i
 import random
 
 # model2
+controlnet1 = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_scribble", torch_dtype=torch.float16)
+pipe1 = StableDiffusionControlNetInpaintPipeline.from_pretrained("hirol/Any-inpainting", controlnet=controlnet1, torch_dtype=torch.float16)
+pipe1.scheduler = UniPCMultistepScheduler.from_config(pipe1.scheduler.config)
+pipe1.to('cuda')
+
+# model1
 controlnet = ControlNetModel.from_pretrained("hirol/control_any5_openpose", torch_dtype=torch.float16)
 pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained("hirol/Any-inpainting", controlnet=controlnet, torch_dtype=torch.float16)
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-
+pipe.to('cuda')
 
 def make_inpaint_condition(image, image_mask):
     image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
@@ -23,16 +29,6 @@ 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)
@@ -56,17 +52,7 @@ 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)
-    # 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,