Update app.py

app.py

@@ -1,431 +1,393 @@
-import os
+import logging
 import random
-import sys
-from typing import Sequence, Mapping, Any, Union
-import torch
+import warnings
+import os
 import gradio as gr
-from huggingface_hub import hf_hub_download
+import numpy as np
 import spaces
-
-# Download required models from Hugging Face
-hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="ae.safetensors", local_dir="models/vae")
-hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="flux1-dev.safetensors", local_dir="models/diffusion_models")
-hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.safetensors", local_dir="models/text_encoders")
-hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir="models/text_encoders")
-hf_hub_download(repo_id="kim2091/UltraSharp", filename="4x-UltraSharp.pth", local_dir="models/upscale_models")
-
-def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    """Returns the value at the given index of a sequence or mapping."""
-    try:
-        return obj[index]
-    except KeyError:
-        return obj["result"][index]
-
-def find_path(name: str, path: str = None) -> str:
-    """Recursively looks at parent folders starting from the given path until it finds the given name."""
-    if path is None:
-        path = os.getcwd()
-
-    if name in os.listdir(path):
-        path_name = os.path.join(path, name)
-        print(f"{name} found: {path_name}")
-        return path_name
-
-    parent_directory = os.path.dirname(path)
-    if parent_directory == path:
-        return None
-
-    return find_path(name, parent_directory)
-
-def add_comfyui_directory_to_sys_path() -> None:
-    """Add 'ComfyUI' to the sys.path"""
-    comfyui_path = find_path("ComfyUI")
-    if comfyui_path is not None and os.path.isdir(comfyui_path):
-        sys.path.append(comfyui_path)
-        print(f"'{comfyui_path}' added to sys.path")
-
-def add_extra_model_paths() -> None:
-    """Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path."""
-    try:
-        from main import load_extra_path_config
-        extra_model_paths = find_path("extra_model_paths.yaml")
-        if extra_model_paths is not None:
-            load_extra_path_config(extra_model_paths)
-        else:
-            print("Could not find the extra_model_paths config file.")
-    except ImportError:
-        try:
-            from utils.extra_config import load_extra_path_config
-            extra_model_paths = find_path("extra_model_paths.yaml")
-            if extra_model_paths is not None:
-                load_extra_path_config(extra_model_paths)
-            else:
-                print("Could not find the extra_model_paths config file.")
-        except ImportError:
-            print("Could not import extra config. Continuing without extra model paths.")
-
-add_comfyui_directory_to_sys_path()
-try:
-    add_extra_model_paths()
-except Exception as e:
-    print(f"Warning: Could not load extra model paths: {e}")
-
-def import_custom_nodes() -> None:
-    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS"""
+import torch
+from diffusers import FluxControlNetModel, FluxControlNetPipeline
+from transformers import AutoProcessor, AutoModelForCausalLM
+from gradio_imageslider import ImageSlider
+from PIL import Image
+from huggingface_hub import snapshot_download
+import requests
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 800px;
+}
+.main-header {
+    text-align: center;
+    margin-bottom: 2rem;
+}
+"""
+
+# Device setup
+if torch.cuda.is_available():
+    power_device = "GPU"
+    device = "cuda"
+else:
+    power_device = "CPU"
+    device = "cpu"
+
+# Get HuggingFace token
+huggingface_token = os.getenv("HF_TOKEN")
+
+# Download FLUX model
+print("📥 Downloading FLUX model...")
+model_path = snapshot_download(
+    repo_id="black-forest-labs/FLUX.1-dev", 
+    repo_type="model", 
+    ignore_patterns=["*.md", "*..gitattributes"],
+    local_dir="FLUX.1-dev",
+    token=huggingface_token,
+)
+
+# Load Florence-2 model for image captioning
+print("📥 Loading Florence-2 model...")
+florence_model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Florence-2-large", 
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+).to(device)
+florence_processor = AutoProcessor.from_pretrained(
+    "microsoft/Florence-2-large", 
+    trust_remote_code=True
+)
+
+# Load FLUX ControlNet pipeline
+print("📥 Loading FLUX ControlNet...")
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", 
+    torch_dtype=torch.bfloat16
+).to(device)
+
+pipe = FluxControlNetPipeline.from_pretrained(
+    model_path, 
+    controlnet=controlnet, 
+    torch_dtype=torch.bfloat16
+)
+pipe.to(device)
+
+print("✅ All models loaded successfully!")
+
+MAX_SEED = 1000000
+MAX_PIXEL_BUDGET = 1024 * 1024
+
+
+def generate_caption(image):
+    """Generate detailed caption using Florence-2"""
     try:
-        import asyncio
-        import execution
-        from nodes import init_extra_nodes
-        import server
-
-        # Check if we're already in an event loop
-        try:
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                # We're in an existing loop, use it
-                pass
-            else:
-                # Loop exists but not running, set a new one
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-        except RuntimeError:
-            # No loop exists, create one
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        server_instance = server.PromptServer(loop)
-        execution.PromptQueue(server_instance)
-        init_extra_nodes()
+        task_prompt = "<MORE_DETAILED_CAPTION>"
+        prompt = task_prompt
+
+        inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to(device)
+        
+        generated_ids = florence_model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            num_beams=3,
+            do_sample=True,
+        )
+        
+        generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        parsed_answer = florence_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+        
+        caption = parsed_answer[task_prompt]
+        return caption
     except Exception as e:
-        print(f"Warning: Could not initialize custom nodes: {e}")
-        print("Continuing with basic ComfyUI nodes only...")
-
-from nodes import NODE_CLASS_MAPPINGS
-
-# Pre-load models outside the decorated function for ZeroGPU efficiency
-try:
-    import_custom_nodes()
-
-    # Initialize model loaders
-    dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-    dualcliploader_54 = dualcliploader.load_clip(
-        clip_name1="clip_l.safetensors",
-        clip_name2="t5xxl_fp16.safetensors",
-        type="flux",
-        device="default",
-    )
+        print(f"Caption generation failed: {e}")
+        return "a high quality detailed image"
 
-    upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
-    upscalemodelloader_44 = upscalemodelloader.load_model(model_name="4x-UltraSharp.pth")
 
-    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-    vaeloader_55 = vaeloader.load_vae(vae_name="ae.safetensors")
+def process_input(input_image, upscale_factor):
+    """Process input image and handle size constraints"""
+    w, h = input_image.size
+    w_original, h_original = w, h
+    aspect_ratio = w / h
 
-    unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-    unetloader_58 = unetloader.load_unet(
-        unet_name="flux1-dev.safetensors", weight_dtype="default"
-    )
+    was_resized = False
 
-    downloadandloadflorence2model = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
-    downloadandloadflorence2model_52 = downloadandloadflorence2model.loadmodel(
-        model="microsoft/Florence-2-large", precision="fp16", attention="sdpa"
-    )
+    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
+        warnings.warn(
+            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to fit budget."
+        )
+        gr.Info(
+            f"Requested output image is too large. Resizing input to fit within pixel budget."
+        )
+        input_image = input_image.resize(
+            (
+                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
+                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
+            )
+        )
+        was_resized = True
+
+    # Resize to multiple of 8
+    w, h = input_image.size
+    w = w - w % 8
+    h = h - h % 8
 
-    # Pre-load models to GPU for efficiency
+    return input_image.resize((w, h)), w_original, h_original, was_resized
+
+
+def load_image_from_url(url):
+    """Load image from URL"""
     try:
-        from comfy import model_management
-        model_loaders = [dualcliploader_54, vaeloader_55, unetloader_58, downloadandloadflorence2model_52]
-        valid_models = [
-            getattr(loader[0], 'patcher', loader[0])
-            for loader in model_loaders
-            if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
-        ]
-        model_management.load_models_gpu(valid_models)
-        print("Models successfully pre-loaded to GPU")
+        response = requests.get(url)
+        response.raise_for_status()
+        return Image.open(requests.get(url, stream=True).raw)
     except Exception as e:
-        print(f"Warning: Could not pre-load models to GPU: {e}")
-
-    print("ComfyUI setup completed successfully!")
+        raise gr.Error(f"Failed to load image from URL: {e}")
+
+
+@spaces.GPU(duration=120)
+def enhance_image(
+    image_input,
+    image_url,
+    seed,
+    randomize_seed,
+    num_inference_steps,
+    upscale_factor,
+    controlnet_conditioning_scale,
+    guidance_scale,
+    use_generated_caption,
+    custom_prompt,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Main enhancement function"""
+    # Handle image input
+    if image_input is not None:
+        input_image = image_input
+    elif image_url:
+        input_image = load_image_from_url(image_url)
+    else:
+        raise gr.Error("Please provide an image (upload or URL)")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    true_input_image = input_image
     
-except Exception as e:
-    print(f"Error during ComfyUI setup: {e}")
-    print("Please check that all required custom nodes are installed.")
-    raise
-
-@spaces.GPU(duration=120)  # Adjust duration based on your workflow speed
-def enhance_image(image_input, upscale_factor, steps, cfg_scale, denoise_strength, guidance_scale):
-    """
-    Main function to enhance and upscale images using Florence-2 captioning and FLUX upscaling
-    """
-    try:
-        with torch.inference_mode():
-            # Handle different input types (file upload vs URL)
-            if isinstance(image_input, str) and image_input.startswith(('http://', 'https://')):
-                # Load from URL
-                load_image_from_url_mtb = NODE_CLASS_MAPPINGS["Load Image From Url (mtb)"]()
-                load_image_result = load_image_from_url_mtb.load(url=image_input)
-            else:
-                # Load from uploaded file
-                loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
-                load_image_result = loadimage.load_image(image=image_input)
-
-            # Generate detailed caption using Florence-2
-            florence2run = NODE_CLASS_MAPPINGS["Florence2Run"]()
-            florence2run_51 = florence2run.encode(
-                text_input="",
-                task="more_detailed_caption",
-                fill_mask=True,
-                keep_model_loaded=False,
-                max_new_tokens=1024,
-                num_beams=3,
-                do_sample=True,
-                output_mask_select="",
-                seed=random.randint(1, 2**64),
-                image=get_value_at_index(load_image_result, 0),
-                florence2_model=get_value_at_index(downloadandloadflorence2model_52, 0),
-            )
+    # Process input image
+    input_image, w_original, h_original, was_resized = process_input(
+        input_image, upscale_factor
+    )
 
-            # Encode the generated caption
-            cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
-            cliptextencode_6 = cliptextencode.encode(
-                text=get_value_at_index(florence2run_51, 2),
-                clip=get_value_at_index(dualcliploader_54, 0),
-            )
+    # Generate caption if requested
+    if use_generated_caption:
+        gr.Info("🔍 Generating image caption...")
+        generated_caption = generate_caption(input_image)
+        prompt = generated_caption
+    else:
+        prompt = custom_prompt if custom_prompt.strip() else ""
 
-            # Encode empty negative prompt
-            cliptextencode_42 = cliptextencode.encode(
-                text="", clip=get_value_at_index(dualcliploader_54, 0)
-            )
+    # Rescale with upscale factor
+    w, h = input_image.size
+    control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
 
-            # Set up upscale factor
-            primitivefloat = NODE_CLASS_MAPPINGS["PrimitiveFloat"]()
-            primitivefloat_60 = primitivefloat.execute(value=upscale_factor)
+    generator = torch.Generator().manual_seed(seed)
 
-            # Apply FLUX guidance
-            fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
-            fluxguidance_26 = fluxguidance.append(
-                guidance=guidance_scale, 
-                conditioning=get_value_at_index(cliptextencode_6, 0)
+    gr.Info("🚀 Upscaling image...")
+    
+    # Generate upscaled image
+    image = pipe(
+        prompt=prompt,
+        control_image=control_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        height=control_image.size[1],
+        width=control_image.size[0],
+        generator=generator,
+    ).images[0]
+
+    if was_resized:
+        gr.Info(f"📏 Resizing output to target size: {w_original * upscale_factor}x{h_original * upscale_factor}")
+
+    # Resize to target desired size
+    final_image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
+    
+    return [true_input_image, final_image, seed, generated_caption if use_generated_caption else ""]
+
+
+# Create Gradio interface
+with gr.Blocks(css=css, title="🎨 AI Image Enhancer - Florence-2 + FLUX") as demo:
+    gr.HTML("""
+    <div class="main-header">
+        <h1>🎨 AI Image Enhancer</h1>
+        <p>Upload an image or provide a URL to enhance it using Florence-2 captioning and FLUX upscaling</p>
+        <p>Currently running on <strong>{}</strong></p>
+    </div>
+    """.format(power_device))
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.HTML("<h3>📤 Input</h3>")
+            
+            with gr.Tabs():
+                with gr.TabItem("📁 Upload Image"):
+                    input_image = gr.Image(
+                        label="Upload Image",
+                        type="pil",
+                        height=300
+                    )
+                
+                with gr.TabItem("🔗 Image URL"):
+                    image_url = gr.Textbox(
+                        label="Image URL",
+                        placeholder="https://example.com/image.jpg",
+                        value="https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg"
+                    )
+            
+            gr.HTML("<h3>🎛️ Caption Settings</h3>")
+            
+            use_generated_caption = gr.Checkbox(
+                label="Use AI-generated caption (Florence-2)",
+                value=True,
+                info="Generate detailed caption automatically"
             )
-
-            # Perform ultimate upscaling
-            ultimatesdupscale = NODE_CLASS_MAPPINGS["UltimateSDUpscale"]()
-            ultimatesdupscale_50 = ultimatesdupscale.upscale(
-                upscale_by=get_value_at_index(primitivefloat_60, 0),
-                seed=random.randint(1, 2**64),
-                steps=steps,
-                cfg=cfg_scale,
-                sampler_name="euler",
-                scheduler="normal",
-                denoise=denoise_strength,
-                mode_type="Linear",
-                tile_width=1024,
-                tile_height=1024,
-                mask_blur=8,
-                tile_padding=32,
-                seam_fix_mode="None",
-                seam_fix_denoise=1,
-                seam_fix_width=64,
-                seam_fix_mask_blur=8,
-                seam_fix_padding=16,
-                force_uniform_tiles=True,
-                tiled_decode=False,
-                image=get_value_at_index(load_image_result, 0),
-                model=get_value_at_index(unetloader_58, 0),
-                positive=get_value_at_index(fluxguidance_26, 0),
-                negative=get_value_at_index(cliptextencode_42, 0),
-                vae=get_value_at_index(vaeloader_55, 0),
-                upscale_model=get_value_at_index(upscalemodelloader_44, 0),
+            
+            custom_prompt = gr.Textbox(
+                label="Custom Prompt (optional)",
+                placeholder="Enter custom prompt or leave empty for generated caption",
+                lines=2
             )
-
-            # Save the result
-            saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
-            saveimage_43 = saveimage.save_images(
-                filename_prefix="enhanced_image",
-                images=get_value_at_index(ultimatesdupscale_50, 0),
+            
+            gr.HTML("<h3>⚙️ Enhancement Settings</h3>")
+            
+            upscale_factor = gr.Slider(
+                label="Upscale Factor",
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=2,
+                info="How much to upscale the image"
             )
-
-            # Return the path to the saved image
-            saved_path = f"output/{saveimage_43['ui']['images'][0]['filename']}"
             
-            # Also return the generated caption for user feedback
-            generated_caption = get_value_at_index(florence2run_51, 2)
+            num_inference_steps = gr.Slider(
+                label="Number of Inference Steps",
+                minimum=8,
+                maximum=50,
+                step=1,
+                value=28,
+                info="More steps = better quality but slower"
+            )
             
-            return saved_path, generated_caption
-
-    except Exception as e:
-        print(f"Error in enhance_image: {str(e)}")
-        raise gr.Error(f"Enhancement failed: {str(e)}")
-
-# Create the Gradio interface
-def create_interface():
-    with gr.Blocks(
-        title="🚀 AI Image Enhancer - Florence-2 + FLUX",
-        theme=gr.themes.Soft(),
-        css="""
-        .gradio-container {
-            max-width: 1200px !important;
-        }
-        .main-header {
-            text-align: center;
-            margin-bottom: 2rem;
-        }
-        .result-gallery {
-            min-height: 400px;
-        }
-        """
-    ) as app:
-        
-        gr.HTML("""
-        <div class="main-header">
-            <h1>🎨 AI Image Enhancer</h1>
-            <p>Upload an image or provide a URL to enhance it using Florence-2 captioning and FLUX upscaling</p>
-        </div>
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.HTML("<h3>📤 Input Settings</h3>")
-                
-                with gr.Tabs():
-                    with gr.TabItem("📁 Upload Image"):
-                        image_upload = gr.Image(
-                            label="Upload Image",
-                            type="filepath",
-                            height=300
-                        )
-                    
-                    with gr.TabItem("🔗 Image URL"):
-                        image_url = gr.Textbox(
-                            label="Image URL",
-                            placeholder="https://example.com/image.jpg",
-                            value="https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg"
-                        )
-                
-                gr.HTML("<h3>⚙️ Enhancement Settings</h3>")
-                
-                upscale_factor = gr.Slider(
-                    minimum=1.0,
-                    maximum=4.0,
-                    value=2.0,
-                    step=0.5,
-                    label="Upscale Factor",
-                    info="How much to upscale the image"
-                )
-                
-                steps = gr.Slider(
-                    minimum=10,
-                    maximum=50,
-                    value=25,
-                    step=5,
-                    label="Steps",
-                    info="Number of denoising steps"
-                )
-                
-                cfg_scale = gr.Slider(
-                    minimum=0.5,
-                    maximum=10.0,
-                    value=1.0,
-                    step=0.5,
-                    label="CFG Scale",
-                    info="Classifier-free guidance scale"
-                )
-                
-                denoise_strength = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.3,
-                    step=0.1,
-                    label="Denoise Strength",
-                    info="How much to denoise the image"
-                )
-                
-                guidance_scale = gr.Slider(
-                    minimum=1.0,
-                    maximum=10.0,
-                    value=3.5,
-                    step=0.5,
-                    label="Guidance Scale",
-                    info="FLUX guidance strength"
+            controlnet_conditioning_scale = gr.Slider(
+                label="ControlNet Conditioning Scale",
+                minimum=0.1,
+                maximum=1.5,
+                step=0.1,
+                value=0.6,
+                info="How much to preserve original structure"
+            )
+            
+            guidance_scale = gr.Slider(
+                label="Guidance Scale",
+                minimum=1.0,
+                maximum=10.0,
+                step=0.5,
+                value=3.5,
+                info="How closely to follow the prompt"
+            )
+            
+            with gr.Row():
+                randomize_seed = gr.Checkbox(
+                    label="Randomize seed",
+                    value=True
                 )
-                
-                enhance_btn = gr.Button(
-                    "🚀 Enhance Image",
-                    variant="primary",
-                    size="lg"
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=42,
+                    interactive=True
                 )
             
-            with gr.Column(scale=1):
-                gr.HTML("<h3>📊 Results</h3>")
-                
-                output_image = gr.Image(
-                    label="Enhanced Image",
-                    type="filepath",
-                    height=400,
-                    interactive=False
-                )
-                
-                generated_caption = gr.Textbox(
-                    label="Generated Caption",
-                    placeholder="The AI-generated caption will appear here...",
-                    lines=3,
+            enhance_btn = gr.Button(
+                "🚀 Enhance Image",
+                variant="primary",
+                size="lg"
+            )
+
+        with gr.Column(scale=1):
+            gr.HTML("<h3>📊 Results</h3>")
+            
+            result_slider = ImageSlider(
+                label="Input / Enhanced",
+                type="pil",
+                interactive=True,
+                height=400
+            )
+            
+            with gr.Row():
+                output_seed = gr.Number(
+                    label="Used Seed",
+                    precision=0,
                     interactive=False
                 )
-                
-                gr.HTML("""
-                <div style="margin-top: 1rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
-                    <h4>💡 How it works:</h4>
-                    <ol>
-                        <li>Florence-2 analyzes your image and generates a detailed caption</li>
-                        <li>FLUX uses this caption to guide the upscaling process</li>
-                        <li>The result is an enhanced, higher-resolution image</li>
-                    </ol>
-                </div>
-                """)
-
-        # Event handlers
-        def process_image(img_upload, img_url, upscale_f, steps_val, cfg_val, denoise_val, guidance_val):
-            # Determine input source
-            image_input = img_upload if img_upload is not None else img_url
             
-            if not image_input:
-                raise gr.Error("Please provide an image (upload or URL)")
-            
-            return enhance_image(image_input, upscale_f, steps_val, cfg_val, denoise_val, guidance_val)
-
-        enhance_btn.click(
-            fn=process_image,
-            inputs=[
-                image_upload,
-                image_url,
-                upscale_factor,
-                steps,
-                cfg_scale,
-                denoise_strength,
-                guidance_scale
-            ],
-            outputs=[output_image, generated_caption]
-        )
-        
-        # Example inputs
-        gr.Examples(
-            examples=[
-                [None, "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg", 2.0, 25, 1.0, 0.3, 3.5],
-                [None, "https://picsum.photos/512/512", 2.0, 20, 1.5, 0.4, 4.0],
-            ],
-            inputs=[
-                image_upload,
-                image_url,
-                upscale_factor,
-                steps,
-                cfg_scale,
-                denoise_strength,
-                guidance_scale
-            ]
-        )
+            generated_caption_output = gr.Textbox(
+                label="Generated Caption",
+                placeholder="AI-generated caption will appear here...",
+                lines=3,
+                interactive=False
+            )
+
+    # Examples
+    gr.Examples(
+        examples=[
+            [None, "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg", 42, False, 28, 2, 0.6, 3.5, True, ""],
+            [None, "https://picsum.photos/512/512", 123, False, 25, 3, 0.8, 4.0, True, ""],
+        ],
+        inputs=[
+            input_image,
+            image_url,
+            seed,
+            randomize_seed,
+            num_inference_steps,
+            upscale_factor,
+            controlnet_conditioning_scale,
+            guidance_scale,
+            use_generated_caption,
+            custom_prompt,
+        ]
+    )
+
+    # Event handler
+    enhance_btn.click(
+        fn=enhance_image,
+        inputs=[
+            input_image,
+            image_url,
+            seed,
+            randomize_seed,
+            num_inference_steps,
+            upscale_factor,
+            controlnet_conditioning_scale,
+            guidance_scale,
+            use_generated_caption,
+            custom_prompt,
+        ],
+        outputs=[result_slider, output_seed, generated_caption_output]
+    )
 
-    return app
+    gr.HTML("""
+    <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
+        <h4>💡 How it works:</h4>
+        <ol>
+            <li><strong>Florence-2</strong> analyzes your image and generates a detailed caption</li>
+            <li><strong>FLUX ControlNet</strong> uses this caption to guide the upscaling process</li>
+            <li>The result is an enhanced, higher-resolution image with improved details</li>
+        </ol>
+        <p><strong>Note:</strong> Due to memory constraints, output is limited to 1024x1024 pixels total budget.</p>
+    </div>
+    """)
 
 if __name__ == "__main__":
-    app = create_interface()
-    app.launch(share=True, server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)