Update demo to use LoRA weights

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 import spaces  # Uncomment if using ZeroGPU
 
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline, DDPMScheduler
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -11,19 +11,28 @@ model_repo_id = "stabilityai/stable-diffusion-2-1-base"
 
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# pipe = pipe.to(device)
+
+pipe = StableDiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.float16).to(device)
+pipe.scheduler = DDPMScheduler.from_pretrained(model_repo_id, subfolder="scheduler")
+
+folder_of_lora_weights = "ID_Booth_LoRA_weights"
+which_checkpoint = "checkpoint-31-6400"
+lora_name = "pytorch_lora_weights.safetensors"
+
 
 backgrounds_list = ["forest", "city street", "beach", "office", "bus", "laboratory", "factory", "construction site", "hospital", "night club", ""]
 poses_list = ["portrait", "side-portrait"]
-id_list = ["ID_0", "ID_1", "ID_2", "ID_3", "ID_4", "ID_5"]
+id_list = ["ID_0", "ID_1", "ID_2"]
 
-gender_dict = {"ID_0": "male"}
+gender_dict = {"ID_0": "male", "ID_1": "male", "ID_2": "female", "ID_2": "male"}
 MAX_SEED = 10000
 image_size = 512
 
 @spaces.GPU  # Uncomment if using ZeroGPU
 def infer(
+    which_id,
     background,
     pose,
     negative_prompt,
@@ -34,16 +43,20 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
     num_images=1
 ):
+    
+    full_lora_weights_path = f"{folder_of_lora_weights}/{which_id}/{which_checkpoint}/{lora_name}"
+    pipe.load_lora_weights(full_lora_weights_path)
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
     id = "ID_0"
-    gender = gender_dict[id]
+    gender = gender_dict[which_id]
 
     # Construct prompt from dropdown selections
-    prompt = f"face {pose.lower()} photo of {gender} {id} person, {background.lower()} background"
+    prompt = f"face {pose.lower()} photo of {gender} sks person, {background.lower()} background"
 
     print(prompt)
     print(negative_prompt)
@@ -144,6 +157,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click],
         fn=infer,
         inputs=[
+            which_id,
             background,
             pose,
             negative_prompt,