Update app.py

app.py

@@ -23,34 +23,27 @@ dl_cn_config(cn_dir)
 dl_tagger_model(tagger_dir)
 dl_lora_model(lora_dir)
 
-def load_model(lora_dir, cn_dir):
-    dtype = torch.float16
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-    controlnet = ControlNetModel.from_pretrained(cn_dir, torch_dtype=dtype, use_safetensors=True)
-
-    pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-        "cagliostrolab/animagine-xl-3.1", controlnet=controlnet, vae=vae, torch_dtype=torch.float16
-    )
-    pipe.enable_model_cpu_offload()
-    pipe.load_lora_weights(lora_dir, weight_name="syoujomannga_line.safetensors")
-    return pipe
 
 @spaces.GPU(duration=120)
-def predict(input_image_path, prompt, negative_prompt, controlnet_scale):
-    pipe = load_model(lora_dir, cn_dir) 
+def predict(self, lora_model, input_image_path, prompt, negative_prompt, controlnet_scale):
+    # LoRAモデルに基づきpipeを取得
+    pipe = self.load_model(lora_model)
     input_image = Image.open(input_image_path)
     base_image = base_generation(input_image.size, (255, 255, 255, 255)).convert("RGB")
     resize_image = resize_image_aspect_ratio(input_image)
     resize_base_image = resize_image_aspect_ratio(base_image)
     generator = torch.manual_seed(0)
     last_time = time.time()
+    
+    # プロンプト生成
     prompt = "masterpiece, best quality, monochrome, greyscale, lineart, white background, star-shaped pupils, " + prompt
     execute_tags = ["realistic", "nose", "asian"]
     prompt = execute_prompt(execute_tags, prompt)
-    prompt = remove_duplicates(prompt)        
+    prompt = remove_duplicates(prompt)
     prompt = remove_color(prompt)
     print(prompt)
 
+    # 画像生成
     output_image = pipe(
         image=resize_base_image,
         control_image=resize_image,
@@ -66,13 +59,43 @@ def predict(input_image_path, prompt, negative_prompt, controlnet_scale):
     output_image = output_image.resize(input_image.size, Image.LANCZOS)
     return output_image
 
+
 class Img2Img:
     def __init__(self):
         self.demo = self.layout()
         self.tagger_model = None
         self.input_image_path = None
         self.bg_removed_image = None
+        self.pipe = None
+        self.current_lora_model = None
+        
+    def load_model(self, lora_model):
+        # 既に正しいpipeがロードされている場合は再利用
+        if self.pipe and self.current_lora_model == lora_model:
+            return self.pipe  # キャッシュされたpipeを返す
+
+        # 新しいpipeの生成
+        dtype = torch.float16
+        vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype)
+        controlnet = ControlNetModel.from_pretrained(cn_dir, torch_dtype=dtype, use_safetensors=True)
+
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            "cagliostrolab/animagine-xl-3.1", controlnet=controlnet, vae=vae, torch_dtype=dtype
+        )
+        self.pipe.enable_model_cpu_offload()
+
+        # LoRAモデルの設定
+        if lora_model == "とりにく風":
+            self.pipe.load_lora_weights(lora_dir, weight_name="tori29umai_line.safetensors")             
+        elif lora_model == "少女漫画風":
+            self.pipe.load_lora_weights(lora_dir, weight_name="syoujomannga_line.safetensors")        
+        elif lora_model == "劇画調風":
+            self.pipe.load_lora_weights(lora_dir, weight_name="gekiga_line.safetensors")           
 
+        # 現在のlora_modelを保存
+        self.current_lora_model = lora_model
+        return self.pipe
+   
     def process_prompt_analysis(self, input_image_path):
         if self.tagger_model is None:
             self.tagger_model = modelLoad(tagger_dir)
@@ -94,6 +117,7 @@ class Img2Img:
         with gr.Blocks(css=css) as demo:
             with gr.Row():
                 with gr.Column():
+                    self.lora_model = gr.Dropdown(label="Image Style",  choices=["プレーン", "とりにく風", "少女漫画風"], value="プレーン")                    
                     self.input_image_path = gr.Image(label="Input image", type='filepath')
                     self.bg_removed_image_path = gr.Image(label="Background Removed Image", type='filepath')
                     
@@ -120,8 +144,8 @@ class Img2Img:
             )
 
             generate_button.click(
-                fn=predict,
-                inputs=[self.bg_removed_image_path, self.prompt, self.negative_prompt, self.controlnet_scale],
+                fn=self.predict,
+                inputs=[self.lora_model, self.bg_removed_image_path, self.prompt, self.negative_prompt, self.controlnet_scale],
                 outputs=self.output_image
             )
         return demo