update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import numpy as np
 import random
 
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline # DiffusionPipeline
 from peft import PeftModel, PeftConfig
 import torch
 
@@ -25,7 +25,7 @@ else:
 # Cache to avoid re-initializing pipelines repeatedly
 model_cache = {}
 
-def load_pipeline(model_id: str):
+def load_pipeline(model_id: str, lora_scale):
     """
     Loads or retrieves a cached DiffusionPipeline.
     
@@ -38,7 +38,7 @@ def load_pipeline(model_id: str):
     if model_id == "macrdel/unico_proj":
         # Use the specified base model for your LoRA adapter.
         base_model = "CompVis/stable-diffusion-v1-4"
-        pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype)
+        pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype)
         # Load the LoRA weights
         pipe.unet = PeftModel.from_pretrained(
             pipe.unet, 
@@ -53,7 +53,9 @@ def load_pipeline(model_id: str):
             torch_dtype=torch_dtype
         )
     else:
-        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype, safety_checker=None).to(device)
+        pipe.unet.load_state_dict({k: lora_scale * v for k, v in pipe.unet.state_dict().items()})
+        pipe.text_encoder.load_state_dict({k: lora_scale * v for k, v in pipe.text_encoder.state_dict().items()})
     
     pipe.to(device)
     model_cache[model_id] = pipe
@@ -76,7 +78,7 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     # Load the pipeline for the chosen model
-    pipe = load_pipeline(model_id)
+    pipe = load_pipeline(model_id, lora_scale)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -84,12 +86,12 @@ def infer(
     generator = torch.Generator(device=device).manual_seed(seed)
 
     # If using the LoRA model, update the LoRA scale if supported.
-    if model_id == "macrdel/unico_proj":
+    # if model_id == "macrdel/unico_proj":
         # This assumes your pipeline's unet has a method to update the LoRA scale.
-        if hasattr(pipe.unet, "set_lora_scale"):
-            pipe.unet.set_lora_scale(lora_scale)
-        else:
-            print("Warning: LoRA scale adjustment method not found on UNet.")
+        # if hasattr(pipe.unet, "set_lora_scale"):
+        #    pipe.unet.set_lora_scale(lora_scale)
+        # else:
+        #    print("Warning: LoRA scale adjustment method not found on UNet.")
 
     image = pipe(
         prompt=prompt,