Update app.py

app.py

@@ -6,22 +6,91 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-from diffusers import FluxImg2ImgPipeline
 from gradio_imageslider import ImageSlider
 from PIL import Image
-from huggingface_hub import snapshot_download
 import requests
-from transformers import T5TokenizerFast
-
-# For ESRGAN (requires pip install basicsr gfpgan)
-try:
-    from basicsr.archs.rrdbnet_arch import RRDBNet
-    from basicsr.utils import img2tensor, tensor2img
-    USE_ESRGAN = True
-except ImportError:
-    USE_ESRGAN = False
-    warnings.warn("basicsr not installed; falling back to LANCZOS interpolation.")
+import sys
+import subprocess
+from huggingface_hub import hf_hub_download
+import tempfile
+
+# Setup ComfyUI and custom nodes
+if not os.path.exists("ComfyUI"):
+    subprocess.run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI"])
+
+custom_nodes_dir = os.path.join("ComfyUI", "custom_nodes")
+os.makedirs(custom_nodes_dir, exist_ok=True)
+
+# Clone UltimateSDUpscaler
+usd_dir = os.path.join(custom_nodes_dir, "ComfyUI_UltimateSDUpscaler")
+if not os.path.exists(usd_dir):
+    subprocess.run(["git", "clone", "https://github.com/ssitu/ComfyUI_UltimateSDUpscaler", usd_dir])
+
+# Clone comfy_mtb
+mtb_dir = os.path.join(custom_nodes_dir, "comfy_mtb")
+if not os.path.exists(mtb_dir):
+    subprocess.run(["git", "clone", "https://github.com/melMass/comfy_mtb", mtb_dir])
+    # Install requirements
+    if os.path.exists(os.path.join(mtb_dir, "requirements.txt")):
+        subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=mtb_dir)
+
+# Clone KJNodes
+kjn_dir = os.path.join(custom_nodes_dir, "ComfyUI-KJNodes")
+if not os.path.exists(kjn_dir):
+    subprocess.run(["git", "clone", "https://github.com/kijai/ComfyUI-KJNodes", kjn_dir])
+    # Install requirements
+    if os.path.exists(os.path.join(kjn_dir, "requirements.txt")):
+        subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=kjn_dir)
+
+# Download models if not present
+comfy_models_dir = os.path.join("ComfyUI", "models")
+os.makedirs(comfy_models_dir, exist_ok=True)
+
+# UNET (Flux FP8)
+unet_dir = os.path.join(comfy_models_dir, "unet")
+os.makedirs(unet_dir, exist_ok=True)
+if not os.path.exists(os.path.join(unet_dir, "flux1-dev-fp8.safetensors")):
+    hf_hub_download(repo_id="Kijai/flux-fp8", filename="flux1-dev-fp8.safetensors", local_dir=unet_dir)
+
+# CLIP models
+clip_dir = os.path.join(comfy_models_dir, "clip")
+os.makedirs(clip_dir, exist_ok=True)
+if not os.path.exists(os.path.join(clip_dir, "clip_l.safetensors")):
+    hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.safetensors", local_dir=clip_dir)
+if not os.path.exists(os.path.join(clip_dir, "t5xxl_fp8_e4m3fn.safetensors")):
+    hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp8_e4m3fn.safetensors", local_dir=clip_dir)
+
+# VAE
+vae_dir = os.path.join(comfy_models_dir, "vae")
+os.makedirs(vae_dir, exist_ok=True)
+if not os.path.exists(os.path.join(vae_dir, "ae.safetensors")):
+    hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="ae.safetensors", subfolder="vae", local_dir=vae_dir)
+
+# Upscale models
+upscale_dir = os.path.join(comfy_models_dir, "upscale_models")
+os.makedirs(upscale_dir, exist_ok=True)
+for model_name in ["RealESRGAN_x2.pth", "RealESRGAN_x4.pth"]:
+    model_path = os.path.join(upscale_dir, model_name)
+    if not os.path.exists(model_path):
+        url = f"https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/{model_name}"
+        with open(model_path, "wb") as f:
+            f.write(requests.get(url).content)
+
+# Add ComfyUI to sys.path
+sys.path.append(os.path.abspath("ComfyUI"))
+
+# Import custom nodes
+from nodes import NODE_CLASS_MAPPINGS, init_custom_nodes
+init_custom_nodes()
+
+# From the provided script
+def get_value_at_index(obj, index):
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
 
+# CSS and constants similar to original
 css = """
 #col-container {
     margin: 0 auto;
@@ -33,127 +102,37 @@ css = """
 }
 """
 
-# Device setup - Default to CPU, let runtime handle GPU
 power_device = "ZeroGPU"
-device = "cpu"
-
-# Get HuggingFace token
-huggingface_token = os.getenv("HF_TOKEN")
-
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 8192 * 8192  # Increased for tiling support
-
+MAX_PIXEL_BUDGET = 8192 * 8192
 
 def make_divisible_by_16(size):
-    """Adjust size to nearest multiple of 16, stretching if necessary"""
     return ((size // 16) * 16) if (size % 16) < 8 else ((size // 16 + 1) * 16)
 
-
 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
 
     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."
-        )
+        gr.Info("Requested output too large. Resizing input.")
         target_input_pixels = MAX_PIXEL_BUDGET / (upscale_factor ** 2)
         scale = (target_input_pixels / (w * h)) ** 0.5
-        new_w = int(w * scale) // 16 * 16  # Ensure divisible by 16 for Flux compatibility
-        new_h = int(h * scale) // 16 * 16
-        if new_w == 0 or new_h == 0:
-            new_w = max(16, new_w)
-            new_h = max(16, new_h)
+        new_w = max(16, int(w * scale) // 16 * 16)
+        new_h = max(16, int(h * scale) // 16 * 16)
         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):
-    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):
-    """Tiled Img2Img to mimic Ultimate SD Upscaler tiling"""
-    w, h = image.size
-    output = image.copy()  # Start with the control image
-
-    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)
-            if tile_h < 16 or tile_w < 16:  # Skip tiny tiles
-                continue
-            tile = image.crop((x, y, x + tile_w, y + tile_h))
-
-            # Force tile to div by 16
-            new_tile_w = make_divisible_by_16(tile_w)
-            new_tile_h = make_divisible_by_16(tile_h)
-            tile = tile.resize((new_tile_w, new_tile_h), resample=Image.LANCZOS)
-
-            # Run Flux on tile
-            gen_tile = pipe(
-                prompt=prompt,
-                image=tile,
-                strength=strength,
-                num_inference_steps=steps,
-                guidance_scale=guidance,
-                height=new_tile_h,
-                width=new_tile_w,
-                generator=generator,
-            ).images[0]
-
-            # Resize gen_tile back to original tile dimensions
-            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)
-                    effective_overlap_x = min(overlap, tile_w)
-                    effective_overlap_y = min(overlap, tile_h)
-                    if x > 0:
-                        for i in range(effective_overlap_x):
-                            for j in range(tile_h):
-                                mask.putpixel((i, j), int(255 * (i / overlap)))
-                    if y > 0:
-                        for i in range(tile_w):
-                            for j in range(effective_overlap_y):
-                                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, (x, y))
-
-    return output
-
+        raise gr.Error(f"Failed to load image: {e}")
 
 @spaces.GPU(duration=120)
 def enhance_image(
@@ -168,139 +147,119 @@ def enhance_image(
     tile_size,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Main enhancement function"""
-    # Lazy loading of models
-    global pipe, esrgan_model
-    if 'pipe' not in globals():
-        try:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            dtype = torch.bfloat16 if device == "cuda" else torch.float32
-
-            print(f"📥 Loading FLUX Img2Img on {device}...")
-            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="tokenizer_2", token=huggingface_token)
-            pipe = FluxImg2ImgPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-schnell",
-                torch_dtype=dtype,
-                low_cpu_mem_usage=True,
-                device_map="balanced",
-                tokenizer_2=tokenizer_2,
-                token=huggingface_token
-            )
-            pipe.enable_vae_tiling()
-            pipe.enable_vae_slicing()
-            if device == "cuda":
-                pipe.reset_device_map()
-                pipe.enable_model_cpu_offload()
-
-            if USE_ESRGAN:
-                esrgan_path = "4x-UltraSharp.pth"
-                if not os.path.exists(esrgan_path):
-                    url = "https://huggingface.co/uwg/upscaler/resolve/main/ESRGAN/4x-UltraSharp.pth"
-                    with open(esrgan_path, "wb") as f:
-                        f.write(requests.get(url).content)
-                esrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-                state_dict = torch.load(esrgan_path)['params_ema']
-                esrgan_model.load_state_dict(state_dict)
-                esrgan_model.eval()
-                esrgan_model.to(device)
-
-            print("✅ Models loaded successfully!")
-        except Exception as e:
-            print(f"Model loading error: {e}, falling back to CPU")
-            device = "cpu"
-            dtype = torch.float32
-            # Reload on CPU if needed
-            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="tokenizer_2", token=huggingface_token)
-            pipe = FluxImg2ImgPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-schnell",
-                torch_dtype=dtype,
-                low_cpu_mem_usage=True,
-                device_map=None,
-                tokenizer_2=tokenizer_2,
-                token=huggingface_token
-            )
-            pipe.enable_vae_tiling()
-            pipe.enable_vae_slicing()
-
-    # 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)")
+    with torch.inference_mode():
+        # Handle input image
+        if image_input is not None:
+            true_input_image = image_input
+        elif image_url:
+            true_input_image = load_image_from_url(image_url)
+        else:
+            raise gr.Error("Provide an image or URL")
+
+        input_image, w_original, h_original, was_resized = process_input(true_input_image, upscale_factor)
+
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+
+        # Prepare ComfyUI input image
+        input_dir = os.path.join("ComfyUI", "input")
+        os.makedirs(input_dir, exist_ok=True)
+        temp_filename = f"input_{random.randint(0, 1000000)}.png"
+        input_path = os.path.join(input_dir, temp_filename)
+        input_image.save(input_path)
+
+        # Nodes
+        load_image_node = NODE_CLASS_MAPPINGS["LoadImage"]()
+        image_loaded = load_image_node.load_image(image=temp_filename)
+        image = get_value_at_index(image_loaded, 0)
+
+        text_multiline = NODE_CLASS_MAPPINGS["Text Multiline"]()
+        text_out = text_multiline.text_multiline(text=custom_prompt if custom_prompt.strip() else "")
+        prompt_text = get_value_at_index(text_out, 0)
+
+        dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+        clip_out = dualcliploader.load_clip(
+            clip_name1="clip_l.safetensors",
+            clip_name2="t5xxl_fp8_e4m3fn.safetensors",
+            type="flux",
+        )
+        clip = get_value_at_index(clip_out, 0)
+
+        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+        conditioning = get_value_at_index(cliptextencode.encode(text=prompt_text, clip=clip), 0)
+
+        fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
+        positive_out = fluxguidance.append(guidance=3.5, conditioning=conditioning)  # Using 3.5 as in original app
+        positive = get_value_at_index(positive_out, 0)
+
+        conditioningzeroout = NODE_CLASS_MAPPINGS["ConditioningZeroOut"]()
+        negative_out = conditioningzeroout.zero_out(conditioning=conditioning)
+        negative = get_value_at_index(negative_out, 0)
+
+        upscale_name = "RealESRGAN_x2.pth" if upscale_factor == 2 else "RealESRGAN_x4.pth"
+        upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
+        upscale_model = get_value_at_index(upscalemodelloader.load_model(model_name=upscale_name), 0)
+
+        vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+        vae = get_value_at_index(vaeloader.load_vae(vae_name="ae.safetensors"), 0)
+
+        unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+        model = get_value_at_index(unetloader.load_unet(unet_name="flux1-dev-fp8.safetensors", weight_dtype="fp8_e4m3fn"), 0)
+
+        ultimatesdupscale = NODE_CLASS_MAPPINGS["UltimateSDUpscale"]()
+        upscale_out = ultimatesdupscale.upscale(
+            upscale_by=float(upscale_factor),
+            seed=seed,
+            steps=num_inference_steps,
+            cfg=1.0,
+            sampler_name="euler",
+            scheduler="normal",
+            denoise=denoising_strength,
+            mode_type="Linear",
+            tile_width=tile_size,
+            tile_height=tile_size,
+            mask_blur=8,
+            tile_padding=32,
+            seam_fix_mode="None",
+            seam_fix_denoise=1.0,
+            seam_fix_width=64,
+            seam_fix_mask_blur=8,
+            seam_fix_padding=16,
+            force_uniform_tiles=True,
+            tiled_decode=False,
+            image=image,
+            model=model,
+            positive=positive,
+            negative=negative,
+            vae=vae,
+            upscale_model=upscale_model,
+        )
+        upscaled_tensor = get_value_at_index(upscale_out, 0)
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+        # Convert to PIL
+        upscaled_img = Image.fromarray((upscaled_tensor[0].cpu().numpy() * 255).astype(np.uint8))
 
-    true_input_image = input_image
-    
-    # Process input image
-    input_image, w_original, h_original, was_resized = process_input(
-        input_image, upscale_factor
-    )
+        target_w, target_h = w_original * upscale_factor, h_original * upscale_factor
+        if upscaled_img.size != (target_w, target_h):
+            upscaled_img = upscaled_img.resize((target_w, target_h), resample=Image.LANCZOS)
 
-    prompt = custom_prompt if custom_prompt.strip() else ""
-
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    gr.Info("🚀 Upscaling image...")
-
-    # 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)
-
-    # Resize control_image to divisible by 16 (stretching)
-    control_w, control_h = control_image.size
-    new_control_w = make_divisible_by_16(control_w)
-    new_control_h = make_divisible_by_16(control_h)
-    if (new_control_w, new_control_h) != (control_w, control_h):
-        control_image = control_image.resize((new_control_w, new_control_h), resample=Image.LANCZOS)
-
-    # Tiled Flux Img2Img for refinement
-    image = tiled_flux_img2img(
-        pipe,
-        prompt,
-        control_image,
-        denoising_strength,
-        num_inference_steps,
-        3.5,  # Updated guidance_scale to match workflow (3.5)
-        generator,
-        tile_size=tile_size,
-        overlap=32
-    )
+        if was_resized:
+            upscaled_img = upscaled_img.resize((target_w, target_h), resample=Image.LANCZOS)
 
-    # Resize back to original target size if stretched
-    target_w, target_h = w_original * upscale_factor, h_original * upscale_factor
-    if image.size != (target_w, target_h):
-        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
+        resized_input = true_input_image.resize(upscaled_img.size, resample=Image.LANCZOS)
 
-    if was_resized:
-        gr.Info(f"📏 Resizing output to target size: {target_w}x{target_h}")
-        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
-    
-    # Resize input image to match output size for slider alignment
-    resized_input = true_input_image.resize(image.size, resample=Image.LANCZOS)
-    
-    # Move back to CPU to release GPU if possible
-    if device == "cuda":
-        pipe.to("cpu")
-        if USE_ESRGAN:
-            esrgan_model.to("cpu")
-    
-    return [resized_input, image]
+        # Cleanup temp file
+        os.remove(input_path)
 
+        return [resized_input, upscaled_img]
 
-# Create Gradio interface
-with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
+# Gradio interface similar to original
+with gr.Blocks(css=css, title="🎨 AI Image Upscaler - Flux FP8") as demo:
     gr.HTML("""
     <div class="main-header">
-        <h1>🎨 AI Image Upscaler</h1>
-        <p>Upload an image or provide a URL to upscale it using FLUX upscaling</p>
-        <p>Currently running on <strong>{}</strong></p>
+        <h1>🎨 AI Image Upscaler - Flux FP8</h1>
+        <p>Upscale images using Flux FP8 with ComfyUI workflow</p>
+        <p>Running on <strong>{}</strong></p>
     </div>
     """.format(power_device))
 
@@ -310,11 +269,7 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
             
             with gr.Tabs():
                 with gr.TabItem("📁 Upload Image"):
-                    input_image = gr.Image(
-                        label="Upload Image",
-                        type="pil",
-                        height=200  # Made smaller
-                    )
+                    input_image = gr.Image(label="Upload Image", type="pil", height=200)
                 
                 with gr.TabItem("🔗 Image URL"):
                     image_url = gr.Textbox(
@@ -338,17 +293,15 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
                 minimum=1,
                 maximum=4,
                 step=1,
-                value=2,
-                info="How much to upscale the image"
+                value=2
             )
             
             num_inference_steps = gr.Slider(
-                label="Number of Inference Steps",
+                label="Inference Steps",
                 minimum=1,
                 maximum=50,
                 step=1,
-                value=4,
-                info="More steps = better quality but slower (default 4 for schnell)"
+                value=25
             )
             
             denoising_strength = gr.Slider(
@@ -356,8 +309,7 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
                 minimum=0.0,
                 maximum=1.0,
                 step=0.05,
-                value=0.3,
-                info="Controls how much the image is transformed"
+                value=0.3
             )
 
             tile_size = gr.Slider(
@@ -365,42 +317,20 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
                 minimum=256,
                 maximum=2048,
                 step=64,
-                value=1024,
-                info="Size of tiles for processing (larger = faster but more memory)"
+                value=1024
             )
             
             with gr.Row():
-                randomize_seed = gr.Checkbox(
-                    label="Randomize seed",
-                    value=True
-                )
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=42,
-                    interactive=True
-                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
             
-            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
+        with gr.Column(scale=2):
             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
-            )
+            result_slider = ImageSlider(type="pil", interactive=False, height=600, label=None)
 
-    # Event handler
     enhance_btn.click(
         fn=enhance_image,
         inputs=[
@@ -419,76 +349,26 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
     
     gr.HTML("""
     <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
-        <p><strong>Note:</strong> This upscaler uses the Flux.1-schnell model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
+        <p><strong>Note:</strong> Uses Flux FP8 model. Ensure compliance with licenses for commercial use.</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;
-            position: absolute;
-            top: 0;
-            left: 0;
-        }
+        #result_slider .slider { width: 100% !important; }
+        #result_slider img { object-fit: contain !important; width: 100% !important; height: auto !important; }
+        #result_slider .gr-button-tool, #result_slider .gr-button-undo, #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; }
     </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'));
-            }
+            if (sliderInput) { sliderInput.value = 50; sliderInput.dispatchEvent(new Event('input')); }
         });
     </script>
     """)