Refactor UI structure and add interactive image tab

app.py

@@ -397,7 +397,6 @@ def cleanup(pipeline, loras = None, embeddings = None):
 
 
 # Gen function
-@spaces.GPU
 def gen_img(
     request: SDReq | SDImg2ImgReq | SDInpaintReq
 ):
@@ -503,6 +502,298 @@ with open("data/images/loras/flux.json", "r") as f:
     loras = json.load(f)
 
 
+# Event functions
+def update_fast_generation(model, fast_generation):
+    if fast_generation:
+        return (
+            gr.update(
+                value=3.5
+            ),
+            gr.update(
+                value=8
+            )
+        )
+
+
+def selected_lora_from_gallery(evt: gr.SelectData):
+    return (
+        gr.update(
+            value=evt.index
+        )
+    )
+
+
+def update_selected_lora(custom_lora):
+    link = custom_lora.split("/")
+    
+    if len(link) == 2:
+        model_card = ModelCard.load(custom_lora)
+        trigger_word = model_card.data.get("instance_prompt", "")
+        image_url = f"""https://huggingface.co/{custom_lora}/resolve/main/{model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)}"""
+        
+        custom_lora_info_css = """
+        <style>
+            .custom-lora-info {
+                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', 'Oxygen', 'Ubuntu', 'Cantarell', 'Fira Sans', 'Droid Sans', 'Helvetica Neue', sans-serif;
+                background: linear-gradient(135deg, #4a90e2, #7b61ff);
+                color: white;
+                padding: 16px;
+                border-radius: 8px;
+                box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+                margin: 16px 0;
+            }
+            .custom-lora-header {
+                font-size: 18px;
+                font-weight: 600;
+                margin-bottom: 12px;
+            }
+            .custom-lora-content {
+                display: flex;
+                align-items: center;
+                background-color: rgba(255, 255, 255, 0.1);
+                border-radius: 6px;
+                padding: 12px;
+            }
+            .custom-lora-image {
+                width: 80px;
+                height: 80px;
+                object-fit: cover;
+                border-radius: 6px;
+                margin-right: 16px;
+            }
+            .custom-lora-text h3 {
+                margin: 0 0 8px 0;
+                font-size: 16px;
+                font-weight: 600;
+            }
+            .custom-lora-text small {
+                font-size: 14px;
+                opacity: 0.9;
+            }
+            .custom-trigger-word {
+                background-color: rgba(255, 255, 255, 0.2);
+                padding: 2px 6px;
+                border-radius: 4px;
+                font-weight: 600;
+            }
+        </style>
+        """
+
+        custom_lora_info_html = f"""
+        <div class="custom-lora-info">
+            <div class="custom-lora-header">Custom LoRA: {custom_lora}</div>
+            <div class="custom-lora-content">
+                <img class="custom-lora-image" src="{image_url}" alt="LoRA preview">
+                <div class="custom-lora-text">
+                    <h3>{link[1].replace("-", " ").replace("_", " ")}</h3>
+                    <small>{"Using: <span class='custom-trigger-word'>"+trigger_word+"</span> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}</small>
+                </div>
+            </div>
+        </div>
+        """
+
+        custom_lora_info_html = f"{custom_lora_info_css}{custom_lora_info_html}"
+
+        return (
+            gr.update( # selected_lora
+                value=custom_lora,
+            ),
+            gr.update( # custom_lora_info
+                value=custom_lora_info_html,
+                visible=True
+            )
+        )
+
+    else:
+        return (
+            gr.update( # selected_lora
+                value=custom_lora,
+            ),
+            gr.update( # custom_lora_info
+                value=custom_lora_info_html if len(link) == 0 else "",
+                visible=False
+            )
+        )
+
+
+def add_to_enabled_loras(model, selected_lora, enabled_loras):
+    lora_data = loras
+    try:
+        selected_lora = int(selected_lora)
+        
+        if 0 <= selected_lora: # is the index of the lora in the gallery
+            lora_info = lora_data[selected_lora]
+            enabled_loras.append({
+                "repo_id": lora_info["repo"],
+                "trigger_word": lora_info["trigger_word"]
+            })
+    except ValueError:
+        link = selected_lora.split("/")
+        if len(link) == 2:
+            model_card = ModelCard.load(selected_lora)
+            trigger_word = model_card.data.get("instance_prompt", "")
+            enabled_loras.append({
+                "repo_id": selected_lora,
+                "trigger_word": trigger_word
+            })
+    
+    return (
+        gr.update( # selected_lora
+            value=""
+        ),
+        gr.update( # custom_lora_info
+            value="",
+            visible=False
+        ),
+        gr.update( # enabled_loras
+            value=enabled_loras
+        )
+    )
+
+
+def update_lora_sliders(enabled_loras):
+    sliders = []
+    remove_buttons = []
+    
+    for lora in enabled_loras:
+        sliders.append(
+            gr.update(
+                label=lora.get("repo_id", ""),
+                info=f"Trigger Word: {lora.get('trigger_word', '')}",
+                visible=True,
+                interactive=True
+            )
+        )
+        remove_buttons.append(
+            gr.update(
+                visible=True,
+                interactive=True
+            )
+        )
+    
+    if len(sliders) < 6:
+        for i in range(len(sliders), 6):
+            sliders.append(
+                gr.update(
+                    visible=False
+                )
+            )
+            remove_buttons.append(
+                gr.update(
+                    visible=False
+                )
+            )
+    
+    return *sliders, *remove_buttons
+
+
+def remove_from_enabled_loras(enabled_loras, index):
+    enabled_loras.pop(index)
+    return (
+        gr.update(
+            value=enabled_loras
+        )
+    )
+
+
+@spaces.GPU
+def generate_image(
+        model, prompt, negative_prompt, fast_generation, enabled_loras,
+        lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5,
+        img2img_image, inpaint_image, canny_image, pose_image, depth_image,
+        img2img_strength, inpaint_strength, canny_strength, pose_strength, depth_strength,
+        resize_mode,
+        scheduler, image_height, image_width, image_num_images_per_prompt,
+        image_num_inference_steps, image_guidance_scale, image_seed,
+        refiner, vae
+    ):
+        base_args = {
+            "model": model,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "fast_generation": fast_generation,
+            "loras": None,
+            "resize_mode": resize_mode,
+            "scheduler": scheduler,
+            "height": int(image_height),
+            "width": int(image_width),
+            "num_images_per_prompt": float(image_num_images_per_prompt),
+            "num_inference_steps": float(image_num_inference_steps),
+            "guidance_scale": float(image_guidance_scale),
+            "seed": int(image_seed),
+            "refiner": refiner,
+            "vae": vae,
+            "controlnet_config": None,
+        }
+        base_args = SDReq(**base_args)
+
+        if len(enabled_loras) > 0:
+            base_args.loras = []
+            for enabled_lora, lora_slider in zip(enabled_loras, [lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5]):
+                if enabled_lora.get("repo_id", None):
+                    base_args.loras.append(
+                        {
+                            "repo_id": enabled_lora["repo_id"],
+                            "weight": lora_slider
+                        }
+                    )
+        
+        image = None
+        mask_image = None
+        strength = None
+        
+        if img2img_image:
+            image = img2img_image
+            strength = float(img2img_strength)
+            
+            base_args = SDImg2ImgReq(
+                **base_args.__dict__,
+                image=image,
+                strength=strength
+            )
+        elif inpaint_image:
+            image = inpaint_image['background'] if not all(pixel == (0, 0, 0) for pixel in list(inpaint_image['background'].getdata())) else None
+            mask_image = inpaint_image['layers'][0] if image else None
+            strength = float(inpaint_strength)
+            
+            base_args = SDInpaintReq(
+                **base_args.__dict__,
+                image=image,
+                mask_image=mask_image,
+                strength=strength
+            )
+        elif any([canny_image, pose_image, depth_image]):
+            base_args.controlnet_config = ControlNetReq(
+                controlnets=[],
+                control_images=[],
+                controlnet_conditioning_scale=[]
+            )
+            
+            if canny_image:
+                base_args.controlnet_config.controlnets.append("canny_fl")
+                base_args.controlnet_config.control_images.append(canny_image)
+                base_args.controlnet_config.controlnet_conditioning_scale.append(float(canny_strength))
+            if pose_image:
+                base_args.controlnet_config.controlnets.append("pose_fl")
+                base_args.controlnet_config.control_images.append(pose_image)
+                base_args.controlnet_config.controlnet_conditioning_scale.append(float(pose_strength))
+            if depth_image:
+                base_args.controlnet_config.controlnets.append("depth_fl")
+                base_args.controlnet_config.control_images.append(depth_image)
+                base_args.controlnet_config.controlnet_conditioning_scale.append(float(depth_strength))
+        else:
+            base_args = SDReq(**base_args.__dict__)
+
+        images = gen_img(base_args)
+        
+        return (
+            gr.update(
+                value=images,
+                interactive=True
+            )
+        )
+
+
 # Main Gradio app
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     # Header