fix pipe

app.py

@@ -11,9 +11,6 @@ import os
 import numpy as np
 from PIL import Image
 import zipfile
-import torch
-
-zero = torch.Tensor([0]).cuda()
 
 path = os.getcwd()
 output_dir = f"{path}/output"
@@ -45,7 +42,7 @@ class webui:
     def __init__(self):
         self.demo = gr.Blocks()
 
-    def undercoat(self, input_image, pos_prompt, neg_prompt, alpha_th, thickness):
+    def undercoat(self, input_image, pos_prompt, neg_prompt, alpha_th, thickness, reference_flg, reference_img):
         org_line_image = input_image
         image = pil2cv(input_image)
         image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
@@ -54,18 +51,19 @@ class webui:
         image[index] = [255, 255, 255, 255]
         input_image = cv2pil(image)
 
-        pipe = get_cn_pipeline()
+        pipe = get_cn_pipeline(reference_flg)
         detectors = get_cn_detector(input_image.resize((1024, 1024), Image.ANTIALIAS))
         
 
-        gen_image = generate(pipe, detectors, pos_prompt, neg_prompt)
+        gen_image = generate(pipe, detectors, pos_prompt, neg_prompt, reference_flg, reference_img)
         color_img, unfinished = process(gen_image.resize((image.shape[1], image.shape[0]), Image.ANTIALIAS) , org_line_image, alpha_th, thickness)
         #color_img = color_img.resize((image.shape[1], image.shape[0]) , Image.ANTIALIAS)
 
 
         output_img = Image.alpha_composite(color_img, org_line_image)
         name = randomname(10)
-        os.makedirs(f"{output_dir}")
+        if not os.path.exists(f"{output_dir}"):
+            os.makedirs(f"{output_dir}")
         os.makedirs(f"{output_dir}/{name}")
         output_img.save(f"{output_dir}/{name}/output_image.png")
         org_line_image.save(f"{output_dir}/{name}/line_image.png")
@@ -84,13 +82,15 @@ class webui:
         with self.demo:
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(type="pil", image_mode="RGBA")
-
+                    input_image = gr.Image(type="pil", image_mode="RGBA", label="lineart")
                     pos_prompt = gr.Textbox(value="1girl, blue hair, pink shirts, bestquality, 4K", max_lines=1000, label="positive prompt")                    
                     neg_prompt = gr.Textbox(value=" (worst quality, low quality:1.2), (lowres:1.2), (bad anatomy:1.2), (greyscale, monochrome:1.4)", max_lines=1000, label="negative prompt")
 
                     alpha_th = gr.Slider(maximum = 255, value=100, label = "alpha threshold")
                     thickness = gr.Number(value=5, label="Thickness of correction area (Odd numbers need to be entered)")
+                
+                    reference_image = gr.Image(type="pil", image_mode="RGB", label="reference_image")
+                    reference_flg = gr.Checkbox(value=True, label="reference_flg")
                     #gr.Slider(maximum = 21, value=3, step=2, label = "Thickness of correction area")
 
                     submit = gr.Button(value="Start")
@@ -101,7 +101,7 @@ class webui:
                         output_file = gr.File()
             submit.click(
                 self.undercoat, 
-                inputs=[input_image, pos_prompt, neg_prompt, alpha_th, thickness], 
+                inputs=[input_image, pos_prompt, neg_prompt, alpha_th, thickness, reference_image, reference_flg], 
                 outputs=[output_0, output_file]
             )
 

sd_model.py

@@ -6,7 +6,7 @@ import spaces
 
 device = "cuda"
 
-def get_cn_pipeline():
+def get_cn_pipeline(reference_flg):
     controlnets = [
         ControlNetModel.from_pretrained("./controlnet/lineart", torch_dtype=torch.float16, use_safetensors=True),
         ControlNetModel.from_pretrained("mattyamonaca/controlnet_line2line_xl", torch_dtype=torch.float16)
@@ -17,13 +17,12 @@ def get_cn_pipeline():
         "cagliostrolab/animagine-xl-3.1", controlnet=controlnets, vae=vae, torch_dtype=torch.float16
     )
 
-    #pipe.enable_model_cpu_offload()
-
-    #if pipe.safety_checker is not None:
-    #    pipe.safety_checker = lambda images, **kwargs: (images, [False])
-    
-    #pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-    #pipe.to(device)
+    if reference_flg == True:
+        pipe.load_ip_adapter(
+            "h94/IP-Adapter",
+            subfolder="sdxl_models",
+            weight_name="ip-adapter-plus_sdxl_vit-h.bin"
+        )
 
     return pipe
 
@@ -51,18 +50,30 @@ def get_cn_detector(image):
     return detectors
 
 @spaces.GPU
-def generate(pipe, detectors, prompt, negative_prompt):
+def generate(pipe, detectors, prompt, negative_prompt, reference_flg=False, reference_img=None):
     pipe.to("cuda")
     default_pos = ""
     default_neg = ""
     prompt = default_pos + prompt 
     negative_prompt = default_neg + negative_prompt 
-    print(type(pipe))
-    image = pipe(
-                prompt=prompt,
-                negative_prompt = negative_prompt,
-                image=detectors,
-                num_inference_steps=50,
-                controlnet_conditioning_scale=[1.0, 0.2],
-            ).images[0]
+    
+
+    if reference_flg==False:
+        image = pipe(
+                    prompt=prompt,
+                    negative_prompt = negative_prompt,
+                    image=detectors,
+                    num_inference_steps=50,
+                    controlnet_conditioning_scale=[1.0, 0.2],
+                ).images[0]
+    else:
+        image = pipe(
+                    prompt=prompt,
+                    negative_prompt = negative_prompt,
+                    image=detectors,
+                    num_inference_steps=50,
+                    controlnet_conditioning_scale=[1.0, 0.2],
+                    ip_adapter_image=reference_img,
+                ).images[0]
+
     return image