Update app.py

app.py

@@ -40,17 +40,23 @@ if torch.cuda.is_available():
     device = "cuda"
 else:
     power_device = "CPU"
-    device = "cpu"# Get HuggingFace token
-huggingface_token = os.getenv("HF_TOKEN")# Download FLUX model
-print(" Downloading FLUX model...")
+    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...")
+)
+
+# 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,
@@ -60,15 +66,21 @@ florence_model = AutoModelForCausalLM.from_pretrained(
 florence_processor = AutoProcessor.from_pretrained(
     "microsoft/Florence-2-large", 
     trust_remote_code=True
-)# Load FLUX Img2Img pipeline
-print(" Loading FLUX Img2Img...")
+)
+
+# Load FLUX Img2Img pipeline
+print("📥 Loading FLUX Img2Img...")
 pipe = FluxImg2ImgPipeline.from_pretrained(
     model_path, 
     torch_dtype=torch.bfloat16
 )
 pipe.to(device)
 pipe.enable_vae_tiling()
-pipe.enable_vae_slicing()print(" All models loaded successfully!")# Download ESRGAN model if using
+pipe.enable_vae_slicing()
+
+print("✅ All models loaded successfully!")
+
+# Download ESRGAN model if using
 if USE_ESRGAN:
     esrgan_path = "4x-UltraSharp.pth"
     if not os.path.exists(esrgan_path):
@@ -79,124 +91,149 @@ if USE_ESRGAN:
     state_dict = torch.load(esrgan_path)['params_ema']
     esrgan_model.load_state_dict(state_dict)
     esrgan_model.eval()
-    esrgan_model.to(device)MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 8192 * 8192  # Increased for tiling supportdef generate_caption(image):
+    esrgan_model.to(device)
+
+MAX_SEED = 1000000
+MAX_PIXEL_BUDGET = 8192 * 8192  # Increased for tiling support
+
+
+def generate_caption(image):
     """Generate detailed caption using Florence-2"""
     try:
         task_prompt = "<MORE_DETAILED_CAPTION>"
-        prompt = task_prompt    inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to(device)
-    inputs["pixel_values"] = inputs["pixel_values"].to(torch.float16)  # Match model dtype
-    
-    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"Caption generation failed: {e}")
-    return "a high quality detailed image"def process_input(input_image, upscale_factor):
+        prompt = task_prompt
+
+        inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to(device)
+        inputs["pixel_values"] = inputs["pixel_values"].to(torch.float16)  # Match model dtype
+        
+        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"Caption generation failed: {e}")
+        return "a high quality detailed image"
+
+
+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 / hwas_resized = False
+    aspect_ratio = w / h
 
-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."
-    )
-    target_input_pixels = MAX_PIXEL_BUDGET / (upscale_factor ** 2)
-    scale = (target_input_pixels / (w * h)) ** 0.5
-    new_w = int(w * scale) - int(w * scale) % 8
-    new_h = int(h * scale) - int(h * scale) % 8
-    input_image = input_image.resize((new_w, new_h), resample=Image.LANCZOS)
-    was_resized = True
-
-return input_image, w_original, h_original, was_resizeddef load_image_from_url(url):
+    was_resized = False
+
+    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."
+        )
+        target_input_pixels = MAX_PIXEL_BUDGET / (upscale_factor ** 2)
+        scale = (target_input_pixels / (w * h)) ** 0.5
+        new_w = int(w * scale) - int(w * scale) % 8
+        new_h = int(h * scale) - int(h * scale) % 8
+        input_image = input_image.resize((new_w, new_h), resample=Image.LANCZOS)
+        was_resized = True
+
+    return input_image, w_original, h_original, was_resized
+
+
+def load_image_from_url(url):
     """Load image from URL"""
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()
         return Image.open(response.raw)
     except Exception as e:
-        raise gr.Error(f"Failed to load image from URL: {e}")def esrgan_upscale(image, scale=4):
+        raise gr.Error(f"Failed to load image from URL: {e}")
+
+
+def esrgan_upscale(image, scale=4):
     if not USE_ESRGAN:
         return image.resize((image.width * scale, image.height * scale), resample=Image.LANCZOS)
     img = img2tensor(np.array(image) / 255., bgr2rgb=False, float32=True)
     with torch.no_grad():
         output = esrgan_model(img.unsqueeze(0)).squeeze()
     output_img = tensor2img(output, rgb2bgr=False, min_max=(0, 1))
-    return Image.fromarray(output_img)def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator, tile_size=1024, overlap=32):
+    return Image.fromarray(output_img)
+
+
+def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator, tile_size=1024, overlap=32):
     """Tiled Img2Img to mimic Ultimate SD Upscaler tiling"""
     w, h = image.size
-    output = image.copy()  # Start with the control image# For handling long prompts: truncate for CLIP, full for T5
-max_clip_tokens = pipe.tokenizer.model_max_length  # Typically 77
-input_ids = pipe.tokenizer.encode(prompt, return_tensors="pt")
-if input_ids.shape[1] > max_clip_tokens:
-    input_ids = input_ids[:, :max_clip_tokens]
-    prompt_clip = pipe.tokenizer.decode(input_ids[0], skip_special_tokens=True)
-else:
-    prompt_clip = prompt
-
-for x in range(0, w, tile_size - overlap):
-    for y in range(0, h, tile_size - overlap):
-        tile_w = min(tile_size, w - x)
-        tile_h = min(tile_size, h - y)
-        tile = image.crop((x, y, x + tile_w, y + tile_h))
-
-        # Run Flux on tile
-        gen_tile = pipe(
-            prompt=prompt_clip,
-            prompt_2=prompt,
-            image=tile,
-            strength=strength,
-            num_inference_steps=steps,
-            guidance_scale=guidance,
-            height=tile_h,
-            width=tile_w,
-            generator=generator,
-        ).images[0]
-
-        # Resize back to exact tile size if pipeline adjusted it
-        gen_tile = gen_tile.resize((tile_w, tile_h), resample=Image.LANCZOS)
-
-        # Paste with blending if overlap
-        if overlap > 0:
-            paste_box = (x, y, x + tile_w, y + tile_h)
-            if x > 0 or y > 0:
-                # Simple linear blend on overlaps
-                mask = Image.new('L', (tile_w, tile_h), 255)
-                if x > 0:
-                    blend_width = min(overlap, tile_w)
-                    for i in range(blend_width):
-                        for j in range(tile_h):
-                            mask.putpixel((i, j), int(255 * (i / overlap)))
-                if y > 0:
-                    blend_height = min(overlap, tile_h)
-                    for i in range(tile_w):
-                        for j in range(blend_height):
-                            mask.putpixel((i, j), int(255 * (j / overlap)))
-                output.paste(gen_tile, paste_box, mask)
+    output = image.copy()  # Start with the control image
+
+    # For handling long prompts: truncate for CLIP, full for T5
+    max_clip_tokens = pipe.tokenizer.model_max_length  # Typically 77
+    input_ids = pipe.tokenizer.encode(prompt, return_tensors="pt")
+    if input_ids.shape[1] > max_clip_tokens:
+        input_ids = input_ids[:, :max_clip_tokens]
+        prompt_clip = pipe.tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    else:
+        prompt_clip = prompt
+
+    for x in range(0, w, tile_size - overlap):
+        for y in range(0, h, tile_size - overlap):
+            tile_w = min(tile_size, w - x)
+            tile_h = min(tile_size, h - y)
+            tile = image.crop((x, y, x + tile_w, y + tile_h))
+
+            # Run Flux on tile
+            gen_tile = pipe(
+                prompt=prompt_clip,
+                prompt_2=prompt,
+                image=tile,
+                strength=strength,
+                num_inference_steps=steps,
+                guidance_scale=guidance,
+                height=tile_h,
+                width=tile_w,
+                generator=generator,
+            ).images[0]
+
+            # Resize back to exact tile size if pipeline adjusted it
+            gen_tile = gen_tile.resize((tile_w, tile_h), resample=Image.LANCZOS)
+
+            # Paste with blending if overlap
+            if overlap > 0:
+                paste_box = (x, y, x + tile_w, y + tile_h)
+                if x > 0 or y > 0:
+                    # Simple linear blend on overlaps
+                    mask = Image.new('L', (tile_w, tile_h), 255)
+                    if x > 0:
+                        blend_width = min(overlap, tile_w)
+                        for i in range(blend_width):
+                            for j in range(tile_h):
+                                mask.putpixel((i, j), int(255 * (i / overlap)))
+                    if y > 0:
+                        blend_height = min(overlap, tile_h)
+                        for i in range(tile_w):
+                            for j in range(blend_height):
+                                mask.putpixel((i, j), int(255 * (j / overlap)))
+                    output.paste(gen_tile, paste_box, mask)
+                else:
+                    output.paste(gen_tile, paste_box)
             else:
-                output.paste(gen_tile, paste_box)
-        else:
-            output.paste(gen_tile, (x, y))
+                output.paste(gen_tile, (x, y))
 
-return output@spaces.GPU(duration=120)
+    return output
+
+
+@spaces.GPU(duration=120)
 def enhance_image(
     image_input,
     image_url,
-    seed,
     randomize_seed,
     num_inference_steps,
     upscale_factor,
@@ -212,244 +249,259 @@ def enhance_image(
     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)
+        raise gr.Error("Please provide an image (upload or URL)")
 
-true_input_image = input_image
+    # Convert input image to PNG in backend
+    buffer = io.BytesIO()
+    input_image.save(buffer, format="PNG")
+    buffer.seek(0)
+    input_image = Image.open(buffer)
 
-# Process input image
-input_image, w_original, h_original, was_resized = process_input(
-    input_image, upscale_factor
-)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    else:
+        seed = 42
 
-# 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 ""
+    true_input_image = input_image
+    
+    # Process input image
+    input_image, w_original, h_original, was_resized = process_input(
+        input_image, upscale_factor
+    )
 
-generator = torch.Generator().manual_seed(seed)
+    # 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 ""
 
-gr.Info(" Upscaling image...")
+    generator = torch.Generator().manual_seed(seed)
 
-# Initial upscale
-if USE_ESRGAN and upscale_factor == 4:
-    control_image = esrgan_upscale(input_image, upscale_factor)
-else:
-    w, h = input_image.size
-    control_image = input_image.resize((w * upscale_factor, h * upscale_factor), resample=Image.LANCZOS)
+    gr.Info("🚀 Upscaling image...")
 
-# Tiled Flux Img2Img for refinement
-image = tiled_flux_img2img(
-    pipe,
-    prompt,
-    control_image,
-    denoising_strength,
-    num_inference_steps,
-    1.0,  # Hardcoded guidance_scale to 1
-    generator,
-    tile_size=1024,
-    overlap=32
-)
+    # Initial upscale
+    if USE_ESRGAN and upscale_factor == 4:
+        control_image = esrgan_upscale(input_image, upscale_factor)
+    else:
+        w, h = input_image.size
+        control_image = input_image.resize((w * upscale_factor, h * upscale_factor), resample=Image.LANCZOS)
+
+    # Tiled Flux Img2Img for refinement
+    image = tiled_flux_img2img(
+        pipe,
+        prompt,
+        control_image,
+        denoising_strength,
+        num_inference_steps,
+        1.0,  # Hardcoded guidance_scale to 1
+        generator,
+        tile_size=1024,
+        overlap=32
+    )
 
-if was_resized:
-    gr.Info(f" Resizing output to target size: {w_original * upscale_factor}x{h_original * upscale_factor}")
-    image = image.resize((w_original * upscale_factor, h_original * upscale_factor), resample=Image.LANCZOS)
+    if was_resized:
+        gr.Info(f"📏 Resizing output to target size: {w_original * upscale_factor}x{h_original * upscale_factor}")
+        image = image.resize((w_original * upscale_factor, h_original * upscale_factor), resample=Image.LANCZOS)
+    
+    # Resize input image to match output size for slider alignment
+    resized_input = true_input_image.resize(image.size, resample=Image.LANCZOS)
+    
+    return [resized_input, image], image
 
-# Resize input image to match output size for slider alignment
-resized_input = true_input_image.resize(image.size, resample=Image.LANCZOS)
 
-return [resized_input, image]# Create Gradio interface
-with gr.Blocks(css=css, title=" AI Image Upscaler - Florence-2 + FLUX") as demo:
+# Create Gradio interface
+with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FLUX") as demo:
     gr.HTML("""
     <div class="main-header">
-        <h1>Flux Dev Ultimate HD Upscaler</h1>
-        <p>Upload an image or provide a URL to upscale it using Florence-2 captioning and FLUX upscaling</p>
+        <h1>🎨 Flux dev Creative Upscaler</h1>
+        <p>Upload an image or provide a URL to upscale it using Florence-2 captioning and FLUX dev with Ultimate SD Upscaler</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=200  # Made smaller
-                )
+    """.format(power_device))
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.HTML("<h3>📤 Input</h3>")
             
-            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"
-        )
-        
-        custom_prompt = gr.Textbox(
-            label="Custom Prompt (optional)",
-            placeholder="Enter custom prompt or leave empty for generated caption",
-            lines=2
-        )
-        
-        gr.HTML("<h3> Upscaling Settings</h3>")
-        
-        upscale_factor = gr.Slider(
-            label="Upscale Factor",
-            minimum=1,
-            maximum=4,
-            step=1,
-            value=2,
-            info="How much to upscale the image"
-        )
-        
-        num_inference_steps = gr.Slider(
-            label="Steps (25 Recommended)",
-            minimum=8,
-            maximum=50,
-            step=1,
-            value=25,
-            info="More steps = better quality but slower"
-        )
-        
-        denoising_strength = gr.Slider(
-            label="Creativity (Denoising value)",
-            minimum=0.0,
-            maximum=1.0,
-            step=0.05,
-            value=0.3,
-            info="More>0.3 = Very Creative, Less<0.1 = More consistent, 0.15-0.3 recommended"
-        )
-        
-        with gr.Row():
-            randomize_seed = gr.Checkbox(
-                label="Randomize seed",
-                value=True
+            with gr.Tabs():
+                with gr.TabItem("📁 Upload Image"):
+                    input_image = gr.Image(
+                        label="Upload Image",
+                        type="pil",
+                        height=200  # Made smaller
+                    )
+                
+                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"
+            )
+            
+            custom_prompt = gr.Textbox(
+                label="Custom Prompt (optional)",
+                placeholder="Enter custom prompt or leave empty for generated caption",
+                lines=2
             )
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
+            
+            gr.HTML("<h3>⚙️ Upscaling Settings</h3>")
+            
+            upscale_factor = gr.Slider(
+                label="Upscale Factor",
+                minimum=1,
+                maximum=4,
                 step=1,
-                value=42,
-                interactive=True
+                value=2,
+                info="How much to upscale the image"
+            )
+            
+            num_inference_steps = gr.Slider(
+                label="Steps (25 Recommended)",
+                minimum=8,
+                maximum=50,
+                step=1,
+                value=25,
+                info="More steps = better quality but slower"
+            )
+            
+            denoising_strength = gr.Slider(
+                label="Creativity (Denoising)",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.05,
+                value=0.3,
+                info="Controls how much the image is transformed"
+            )
+            
+            with gr.Row():
+                randomize_seed = gr.Checkbox(
+                    label="Randomize seed",
+                    value=True
+                )
+            
+            enhance_btn = gr.Button(
+                "🚀 Upscale Image",
+                variant="primary",
+                size="lg"
             )
-        
-        enhance_btn = gr.Button(
-            " Upscale Image",
-            variant="primary",
-            size="lg"
-        )
 
-    with gr.Column(scale=2):  # Larger scale for results
-        gr.HTML("<h3> Results</h3>")
-        
-        result_slider = ImageSlider(
-            type="pil",
-            interactive=False,  # Disable interactivity to prevent uploads
-            height=600,  # Made larger
-            elem_id="result_slider",
-            label=None  # Remove default label
-        )
+        with gr.Column(scale=2):  # Larger scale for results
+            gr.HTML("<h3>📊 Results</h3>")
+            
+            result_slider = ImageSlider(
+                type="pil",
+                interactive=False,  # Disable interactivity to prevent uploads
+                height=600,  # Made larger
+                elem_id="result_slider",
+                label=None  # Remove default label
+            )
 
-# Event handler
-enhance_btn.click(
-    fn=enhance_image,
-    inputs=[
-        input_image,
-        image_url,
-        seed,
-        randomize_seed,
-        num_inference_steps,
-        upscale_factor,
-        denoising_strength,
-        use_generated_caption,
-        custom_prompt,
-    ],
-    outputs=[result_slider]
-)
+            upscaled_output = gr.Image(
+                label="Upscaled Image (Download as PNG)",
+                type="pil",
+                interactive=False,
+                show_download_button=True,
+                height=600,
+            )
 
-gr.HTML("""
-<div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
-    <p><strong>Note:</strong> This upscaler uses the Flux dev model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
-</div>
-""")
-
-# Custom CSS for slider
-gr.HTML("""
-<style>
-    #result_slider .slider {
-        width: 100% !important;
-        max-width: inherit !important;
-    }
-    #result_slider img {
-        object-fit: contain !important;
-        width: 100% !important;
-        height: auto !important;
-    }
-    #result_slider .gr-button-tool {
-        display: none !important;
-    }
-    #result_slider .gr-button-undo {
-        display: none !important;
-    }
-    #result_slider .gr-button-clear {
-        display: none !important;
-    }
-    #result_slider .badge-container .badge {
-        display: none !important;
-    }
-    #result_slider .badge-container::before {
-        content: "Before";
-        position: absolute;
-        top: 10px;
-        left: 10px;
-        background: rgba(0,0,0,0.5);
-        color: white;
-        padding: 5px;
-        border-radius: 5px;
-        z-index: 10;
-    }
-    #result_slider .badge-container::after {
-        content: "After";
-        position: absolute;
-        top: 10px;
-        right: 10px;
-        background: rgba(0,0,0,0.5);
-        color: white;
-        padding: 5px;
-        border-radius: 5px;
-        z-index: 10;
-    }
-    #result_slider .fullscreen img {
-        object-fit: contain !important;
-        width: 100vw !important;
-        height: 100vh !important;
-    }
-</style>
-""")
-
-# JS to set slider default position to middle
-gr.HTML("""
-<script>
-    document.addEventListener('DOMContentLoaded', function() {
-        const sliderInput = document.querySelector('#result_slider input[type="range"]');
-        if (sliderInput) {
-            sliderInput.value = 50;
-            sliderInput.dispatchEvent(new Event('input'));
+    # Event handler
+    enhance_btn.click(
+        fn=enhance_image,
+        inputs=[
+            input_image,
+            image_url,
+            randomize_seed,
+            num_inference_steps,
+            upscale_factor,
+            denoising_strength,
+            use_generated_caption,
+            custom_prompt,
+        ],
+        outputs=[result_slider, upscaled_output]
+    )
+    
+    gr.HTML("""
+    <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
+        <p><strong>Note:</strong> This upscaler uses the Flux dev model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
+    </div>
+    """)
+    
+    # Custom CSS for slider
+    gr.HTML("""
+    <style>
+        #result_slider .slider {
+            width: 100% !important;
+            max-width: inherit !important;
+        }
+        #result_slider img {
+            object-fit: contain !important;
+            width: 100% !important;
+            height: auto !important;
+        }
+        #result_slider .gr-button-tool {
+            display: none !important;
+        }
+        #result_slider .gr-button-undo {
+            display: none !important;
         }
-    });
-</script>
-""")if __name__ == "__main__":
-    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)
+        #result_slider .gr-button-clear {
+            display: none !important;
+        }
+        #result_slider .badge-container .badge {
+            display: none !important;
+        }
+        #result_slider .badge-container::before {
+            content: "Before";
+            position: absolute;
+            top: 10px;
+            left: 10px;
+            background: rgba(0,0,0,0.5);
+            color: white;
+            padding: 5px;
+            border-radius: 5px;
+            z-index: 10;
+        }
+        #result_slider .badge-container::after {
+            content: "After";
+            position: absolute;
+            top: 10px;
+            right: 10px;
+            background: rgba(0,0,0,0.5);
+            color: white;
+            padding: 5px;
+            border-radius: 5px;
+            z-index: 10;
+        }
+        #result_slider .fullscreen img {
+            object-fit: contain !important;
+            width: 100vw !important;
+            height: 100vh !important;
+        }
+    </style>
+    """)
+    
+    # JS to set slider default position to middle
+    gr.HTML("""
+    <script>
+        document.addEventListener('DOMContentLoaded', function() {
+            const sliderInput = document.querySelector('#result_slider input[type="range"]');
+            if (sliderInput) {
+                sliderInput.value = 50;
+                sliderInput.dispatchEvent(new Event('input'));
+            }
+        });
+    </script>
+    """)
 
+if __name__ == "__main__":
+    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)