Add LORA weight customization

app.py

@@ -160,7 +160,8 @@ elif model_type == "LORA":
     else:
         raise ValueError(f"Unsupported base model: {base_model}")
 
-    pipe.load_lora_weights(get_file_name("Model"), adapter_name=slugify(model_name))
+    adapter_name = slugify(model_name)
+    pipe.load_lora_weights(get_file_name("Model"), adapter_name=adapter_name)
 else:
     raise ValueError(f"Unsupported model type: {model_type}")
 
@@ -176,6 +177,7 @@ def infer(
     strength: float,
     num_inference_steps: int,
     guidance_scale: float,
+    lora_weight: float,
     progress=gr.Progress(track_tqdm=True),
 ):
     logger.info(f"Starting image generation: {dict(prompt=prompt, image=init_image)}")
@@ -193,6 +195,9 @@ def infer(
         if v
     }
 
+    if lora_weight:
+        pipe.set_adapters(adapter_name, lora_weight)
+
     logger.debug(f"Generating image: {dict(prompt=prompt, **additional_args)}")
 
     images = pipe(
@@ -249,6 +254,15 @@ with gr.Blocks(css=css) as demo:
                     value=0.0,
                 )
 
+                lora_weight = gr.Slider(
+                    label="LORA weight",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.0,
+                    visible=model_type == "LORA",
+                )
+
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=0,
@@ -274,6 +288,7 @@ with gr.Blocks(css=css) as demo:
             strength,
             num_inference_steps,
             guidance_scale,
+            lora_weight,
         ],
         outputs=[result],
     )