Update app.py

app.py

@@ -22,6 +22,7 @@ MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
+ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -47,6 +48,7 @@ def generate(
     guidance_scale_refiner: float = 5.0,
     num_inference_steps_base: int = 25,
     num_inference_steps_refiner: int = 25,
+    use_lora: bool = False,
     apply_refiner: bool = False,
     model = 'stabilityai/stable-diffusion-xl-base-1.0',
     vaecall = 'madebyollin/sdxl-vae-fp16-fix',
@@ -54,6 +56,9 @@ def generate(
 ) -> PIL.Image.Image:
     if torch.cuda.is_available():
         pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+
+        if use_lora:
+            pipe.load_lora_weights(lora)
                                    
         if ENABLE_CPU_OFFLOAD:
             pipe.enable_model_cpu_offload()
@@ -185,6 +190,7 @@ with gr.Blocks(css="style.css") as demo:
                 step=32,
                 value=1024,
             )
+        use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
         apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
         with gr.Row():
             guidance_scale_base = gr.Slider(
@@ -246,6 +252,13 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    use_lora.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_lora,
+        outputs=lora,
+        queue=False,
+        api_name=False,
+    )
     apply_refiner.change(
         fn=lambda x: gr.update(visible=x),
         inputs=apply_refiner,
@@ -284,6 +297,7 @@ with gr.Blocks(css="style.css") as demo:
             guidance_scale_refiner,
             num_inference_steps_base,
             num_inference_steps_refiner,
+            use_lora,
             apply_refiner,
             model,
             vaecall,