Update app.py

app.py

@@ -12,8 +12,6 @@ from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
 import requests
-import io
-import base64
 
 # For ESRGAN (requires pip install basicsr gfpgan)
 try:
@@ -62,7 +60,7 @@ florence_model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Florence-2-large", 
     torch_dtype=torch.float16,
     trust_remote_code=True,
-    attn_implementation="eager"
+    attn_implementation="eager"  # Fix for SDPA compatibility issue
 ).to(device)
 florence_processor = AutoProcessor.from_pretrained(
     "microsoft/Florence-2-large", 
@@ -95,15 +93,16 @@ if USE_ESRGAN:
     esrgan_model.to(device)
 
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 8192 * 8192
+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)
         
         generated_ids = florence_model.generate(
             input_ids=inputs["input_ids"],
@@ -122,10 +121,13 @@ def generate_caption(image):
         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 / h
+
     was_resized = False
 
     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
@@ -144,19 +146,17 @@ def process_input(input_image, upscale_factor):
 
     return input_image, w_original, h_original, was_resized
 
+
 def load_image_from_url(url):
-    """Load image from URL and convert to PNG"""
+    """Load image from URL"""
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()
-        img = Image.open(response.raw)
-        buffer = io.BytesIO()
-        img.save(buffer, format="PNG")
-        buffer.seek(0)
-        return Image.open(buffer)
+        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)
@@ -166,18 +166,11 @@ def esrgan_upscale(image, scale=4):
     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()
-
-    max_clip_tokens = pipe.tokenizer.model_max_length
-    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
+    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):
@@ -185,9 +178,9 @@ def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator
             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,
+                prompt=prompt,
                 image=tile,
                 strength=strength,
                 num_inference_steps=steps,
@@ -197,21 +190,19 @@ def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator
                 generator=generator,
             ).images[0]
 
-            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 i in range(overlap):
                             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):
+                            for j in range(overlap):
                                 mask.putpixel((i, j), int(255 * (j / overlap)))
                     output.paste(gen_tile, paste_box, mask)
                 else:
@@ -221,19 +212,12 @@ def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator
 
     return output
 
-def download_png(image):
-    """Convert image to PNG and return base64 string for download"""
-    if image is None:
-        raise gr.Error("No upscaled image available to download")
-    buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
-    base64_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
-    return base64_data
 
 @spaces.GPU(duration=120)
 def enhance_image(
     image_input,
     image_url,
+    seed,
     randomize_seed,
     num_inference_steps,
     upscale_factor,
@@ -243,11 +227,9 @@ def enhance_image(
     progress=gr.Progress(track_tqdm=True),
 ):
     """Main enhancement function"""
+    # Handle image input
     if image_input is not None:
-        buffer = io.BytesIO()
-        image_input.save(buffer, format="PNG")
-        buffer.seek(0)
-        input_image = Image.open(buffer)
+        input_image = image_input
     elif image_url:
         input_image = load_image_from_url(image_url)
     else:
@@ -255,15 +237,15 @@ def enhance_image(
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    else:
-        seed = 42
 
     true_input_image = input_image
     
+    # Process input image
     input_image, w_original, h_original, was_resized = process_input(
         input_image, upscale_factor
     )
 
+    # Generate caption if requested
     if use_generated_caption:
         gr.Info("🔍 Generating image caption...")
         generated_caption = generate_caption(input_image)
@@ -275,19 +257,21 @@ def enhance_image(
 
     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)
 
+    # Tiled Flux Img2Img for refinement
     image = tiled_flux_img2img(
         pipe,
         prompt,
         control_image,
         denoising_strength,
         num_inference_steps,
-        1.0,
+        1.0,  # Hardcoded guidance_scale to 1
         generator,
         tile_size=1024,
         overlap=32
@@ -297,16 +281,18 @@ def enhance_image(
         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
+    return [resized_input, image]
+
 
 # Create Gradio interface
-with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FLUX") as demo:
+with gr.Blocks(css=css, title="🎨 AI Image Upscaler - Florence-2 + FLUX") as demo:
     gr.HTML("""
     <div class="main-header">
-        <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>
+        <h1>🎨 AI Image Upscaler</h1>
+        <p>Upload an image or provide a URL to upscale it using Florence-2 captioning and FLUX upscaling</p>
         <p>Currently running on <strong>{}</strong></p>
     </div>
     """.format(power_device))
@@ -320,7 +306,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
                     input_image = gr.Image(
                         label="Upload Image",
                         type="pil",
-                        height=200
+                        height=200  # Made smaller
                     )
                 
                 with gr.TabItem("🔗 Image URL"):
@@ -356,7 +342,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
             )
             
             num_inference_steps = gr.Slider(
-                label="Steps (25 Recommended)",
+                label="Number of Inference Steps",
                 minimum=8,
                 maximum=50,
                 step=1,
@@ -365,7 +351,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
             )
             
             denoising_strength = gr.Slider(
-                label="Creativity (Denoising)",
+                label="Denoising Strength",
                 minimum=0.0,
                 maximum=1.0,
                 step=0.05,
@@ -378,6 +364,14 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
                     label="Randomize seed",
                     value=True
                 )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=42,
+                    interactive=True
+                )
             
             enhance_btn = gr.Button(
                 "🚀 Upscale Image",
@@ -385,35 +379,24 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
                 size="lg"
             )
 
-        with gr.Column(scale=2):
+        with gr.Column(scale=2):  # Larger scale for results
             gr.HTML("<h3>📊 Results</h3>")
             
             result_slider = ImageSlider(
                 type="pil",
-                interactive=False,
-                height=600,
+                interactive=False,  # Disable interactivity to prevent uploads
+                height=600,  # Made larger
                 elem_id="result_slider",
-                label=None
-            )
-            
-            download_btn = gr.Button(
-                "📥 Download as PNG",
-                variant="secondary",
-                size="lg"
+                label=None  # Remove default label
             )
 
-    # State to store the upscaled image
-    upscaled_image_state = gr.State()
-
-    # Hidden textbox for base64 data
-    download_data = gr.Textbox(visible=False, elem_id="download_data")
-
-    # Event handlers
+    # Event handler
     enhance_btn.click(
         fn=enhance_image,
         inputs=[
             input_image,
             image_url,
+            seed,
             randomize_seed,
             num_inference_steps,
             upscale_factor,
@@ -421,13 +404,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
             use_generated_caption,
             custom_prompt,
         ],
-        outputs=[result_slider, upscaled_image_state]
-    )
-    
-    download_btn.click(
-        fn=download_png,
-        inputs=[upscaled_image_state],
-        outputs=download_data
+        outputs=[result_slider]
     )
     
     gr.HTML("""
@@ -436,6 +413,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
     </div>
     """)
     
+    # Custom CSS for slider
     gr.HTML("""
     <style>
         #result_slider .slider {
@@ -489,6 +467,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
     </style>
     """)
     
+    # JS to set slider default position to middle
     gr.HTML("""
     <script>
         document.addEventListener('DOMContentLoaded', function() {
@@ -497,31 +476,6 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
                 sliderInput.value = 50;
                 sliderInput.dispatchEvent(new Event('input'));
             }
-
-            const downloadData = document.querySelector('#download_data textarea');
-            if (downloadData) {
-                const observer = new MutationObserver(() => {
-                    const base64 = downloadData.value;
-                    if (base64) {
-                        const byteCharacters = atob(base64);
-                        const byteNumbers = new Array(byteCharacters.length);
-                        for (let i = 0; i < byteCharacters.length; i++) {
-                            byteNumbers[i] = byteCharacters.charCodeAt(i);
-                        }
-                        const byteArray = new Uint8Array(byteNumbers);
-                        const blob = new Blob([byteArray], {type: 'image/png'});
-                        const url = URL.createObjectURL(blob);
-                        const a = document.createElement('a');
-                        a.href = url;
-                        a.download = 'upscaled_image.png';
-                        a.click();
-                        URL.revokeObjectURL(url);
-                        // Clear the textbox
-                        downloadData.value = '';
-                    }
-                });
-                observer.observe(downloadData, {childList: true, subtree: true, characterData: true});
-            }
         });
     </script>
     """)