Update app.py

app.py

@@ -1,12 +1,44 @@
 import gradio as gr
 import numpy as np
 import random
+import os
 
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionPipeline
+from peft import PeftModel, LoraConfig
 import torch
 from typing import Optional
 
+
+def get_lora_sd_pipeline(
+    ckpt_dir='./lora_logos', 
+    base_model_name_or_path=None, 
+    dtype=torch.float16, 
+    adapter_name="default"
+):
+    unet_sub_dir = os.path.join(ckpt_dir, "unet")
+    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
+    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
+        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
+        base_model_name_or_path = config.base_model_name_or_path
+
+    if base_model_name_or_path is None:
+        raise ValueError("Please specify the base model name or path")
+
+    pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype).to(device)
+    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
+
+    if os.path.exists(text_encoder_sub_dir):
+        pipe.text_encoder = PeftModel.from_pretrained(
+            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
+        )
+
+    if dtype in (torch.float16, torch.bfloat16):
+        pipe.unet.half()
+        pipe.text_encoder.half()
+
+    return pipe
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id_default = "CompVis/stable-diffusion-v1-4"
 
@@ -15,7 +47,13 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
-pipe_default = DiffusionPipeline.from_pretrained(model_id_default, torch_dtype=torch_dtype)
+
+pipe_default = get_lora_sd_pipeline(
+    ckpt_dir='./lora_logos', 
+    base_model_name_or_path=model_id_default, 
+    dtype=torch_dtype,
+    )
+# pipe_default = DiffusionPipeline.from_pretrained(model_id_default, torch_dtype=torch_dtype)
 pipe_default = pipe_default.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -32,6 +70,7 @@ def infer(
     model_id: Optional[str] = 'CompVis/stable-diffusion-v1-4',
     seed: Optional[int] = 42,
     guidance_scale: Optional[float] = 7.0,
+    lora_scale: Optional[float] = 0.5,
     progress=gr.Progress(track_tqdm=True),
 ):
     generator = torch.Generator().manual_seed(seed)
@@ -49,8 +88,10 @@ def infer(
     if model_id != model_id_default:
         pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
         pipe = pipe.to(device)
+        pipe.fuse_lora(lora_scale=0.4)
         image = pipe(**params).images[0]
     else:
+        pipe_default.fuse_lora(lora_scale=0.4)
         image = pipe_default(**params).images[0]
 
     return image
@@ -105,6 +146,15 @@ with gr.Blocks(css=css) as demo:
                 value=7.0,
             )
 
+        with gr.Row():
+            lora_scale = gr.Slider(
+                label="LoRA scale",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=0.5,
+            )
+
         with gr.Row():
             num_inference_steps = gr.Slider(
                 label="Number of inference steps",
@@ -148,6 +198,7 @@ with gr.Blocks(css=css) as demo:
             model_id,
             seed,
             guidance_scale,
+            lora_scale,
         ],
         outputs=[result],
     )