Update app.py

app.py

@@ -180,10 +180,19 @@ def prepare_image(image, max_size=MAX_INPUT_SIZE):
     return image
 
 
-def esrgan_upscale(image, model, device='cuda'):
-    """Upscale image 4x using ESRGAN"""
+def esrgan_upscale(image, model, device='cuda', upscale_factor=4):
+    """Upscale image using ESRGAN with variable factor support"""
+    orig_w, orig_h = image.size
+    pre_resize_factor = upscale_factor / 4.0
+    low_res_w = int(orig_w * pre_resize_factor)
+    low_res_h = int(orig_h * pre_resize_factor)
+    if low_res_w < 1 or low_res_h < 1:
+        raise ValueError("Upscale factor too small for image size")
+    
+    low_res_image = image.resize((low_res_w, low_res_h), Image.LANCZOS)
+    
     # Prepare image
-    img_np = np.array(image).astype(np.float32) / 255.
+    img_np = np.array(low_res_image).astype(np.float32) / 255.
     img_np = np.transpose(img_np, (2, 0, 1))  # HWC to CHW
     img_tensor = torch.from_numpy(img_np).unsqueeze(0).to(device)
     
@@ -195,10 +204,18 @@ def esrgan_upscale(image, model, device='cuda'):
         output_np = np.transpose(output_np, (1, 2, 0))  # CHW to HWC
         output_np = (output_np * 255).astype(np.uint8)
     
-    return Image.fromarray(output_np)
+    upscaled = Image.fromarray(output_np)
+    
+    # Resize back to exact target size if needed (due to rounding)
+    target_w = int(orig_w * upscale_factor)
+    target_h = int(orig_h * upscale_factor)
+    if upscaled.size != (target_w, target_h):
+        upscaled = upscaled.resize((target_w, target_h), Image.LANCZOS)
+    
+    return upscaled
 
 
-@spaces.GPU(duration=60)  # 60 seconds should be enough
+@spaces.GPU(duration=120)  # Increased to 120 seconds
 def enhance_image(
     input_image,
     prompt,
@@ -206,6 +223,7 @@ def enhance_image(
     randomize_seed,
     num_inference_steps,
     denoising_strength,
+    upscale_factor,
     progress=gr.Progress(track_tqdm=True),
 ):
     """Main enhancement function"""
@@ -228,12 +246,12 @@ def enhance_image(
         input_image = prepare_image(input_image)
         original_size = input_image.size
         
-        # Step 1: ESRGAN upscale (4x) on GPU
-        gr.Info("🔍 Upscaling with ESRGAN 4x...")
+        # Step 1: ESRGAN upscale on GPU
+        gr.Info(f"🔍 Upscaling with ESRGAN x{upscale_factor}...")
         
         # Move ESRGAN to GPU for faster processing
         esrgan_model.to("cuda")
-        upscaled_image = esrgan_upscale(input_image, esrgan_model, device="cuda")
+        upscaled_image = esrgan_upscale(input_image, esrgan_model, device="cuda", upscale_factor=upscale_factor)
         
         # Move ESRGAN back to CPU to free memory
         esrgan_model.to("cpu")
@@ -299,9 +317,9 @@ def enhance_image(
 with gr.Blocks(css=css) as demo:
     gr.HTML("""
     <div class="main-header">
-        <h1>🚀 ESRGAN 4x + FLUX Enhancement</h1>
-        <p>Upload an image to upscale 4x with ESRGAN and enhance with FLUX</p>
-        <p>Optimized for <strong>ZeroGPU</strong> | Max input: 512x512 → Output: 2048x2048</p>
+        <h1>🚀 Flux Dev Ultimate Upscaler</h1>
+        <p>Upload an image to upscale 2-4x with ESRGAN and enhance with FLUX</p>
+        <p>Optimized for <strong>ZeroGPU</strong> | Max input: 512x512 → Output: up to 2048x2048</p>
     </div>
     """)
     
@@ -315,46 +333,52 @@ with gr.Blocks(css=css) as demo:
             )
             
             prompt = gr.Textbox(
-                label="Enhancement Prompt",
+                label="Describe image with prompt",
                 placeholder="Describe the desired enhancement (e.g., 'high quality, sharp details, vibrant colors')",
                 value="high quality, ultra detailed, sharp",
                 lines=2
             )
             
-            with gr.Accordion("Advanced Settings", open=False):
-                num_inference_steps = gr.Slider(
-                    label="Enhancement Steps",
-                    minimum=10,
-                    maximum=25,
-                    step=1,
-                    value=18,
-                    info="More steps = better quality but slower"
+            # Advanced Settings (always open)
+            upscale_factor = gr.Slider(
+                label="Upscale Ratio",
+                minimum=2,
+                maximum=4,
+                step=1,
+                value=4,
+                info="Choose upscale factor (2x, 3x, 4x)"
+            )
+            
+            num_inference_steps = gr.Slider(
+                label="Enhancement Steps",
+                minimum=10,
+                maximum=25,
+                step=1,
+                value=18,
+                info="More steps = better quality but slower"
+            )
+            
+            denoising_strength = gr.Slider(
+                label="Creativity (Denoising)",
+                minimum=0.1,
+                maximum=0.6,
+                step=0.05,
+                value=0.35,
+                info="Higher = more changes to the image"
+            )
+            
+            with gr.Row():
+                randomize_seed = gr.Checkbox(
+                    label="Randomize seed",
+                    value=True
                 )
-                
-                denoising_strength = gr.Slider(
-                    label="Enhancement Strength",
-                    minimum=0.1,
-                    maximum=0.6,
-                    step=0.05,
-                    value=0.35,
-                    info="Higher = more changes to the image"
+                seed = gr.Number(
+                    label="Seed",
+                    value=42
                 )
-                
-                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
-                    )
             
             enhance_btn = gr.Button(
-                "🎨 Enhance Image (4x Upscale)",
+                "Upscale",
                 variant="primary",
                 size="lg"
             )
@@ -374,18 +398,6 @@ with gr.Blocks(css=css) as demo:
                 visible=False
             )
     
-    # Examples
-    gr.Examples(
-        examples=[
-            ["high quality, ultra detailed, sharp"],
-            ["cinematic, professional photography, enhanced details"],
-            ["vibrant colors, high contrast, sharp focus"],
-            ["photorealistic, 8k quality, fine details"],
-        ],
-        inputs=[prompt],
-        label="Example Prompts"
-    )
-    
     # Event handler
     enhance_btn.click(
         fn=enhance_image,
@@ -396,14 +408,16 @@ with gr.Blocks(css=css) as demo:
             randomize_seed,
             num_inference_steps,
             denoising_strength,
+            upscale_factor,
         ],
         outputs=[result_slider, used_seed]
     )
     
     gr.HTML("""
     <div style="margin-top: 2rem; text-align: center; color: #666;">
-        <p>📌 Pipeline: ESRGAN 4x-UltraSharp → FLUX Dev Enhancement</p>
+        <p>📌 Pipeline: ESRGAN 2-4x-UltraSharp → FLUX Dev Enhancement</p>
         <p>⚡ Optimized for ZeroGPU with automatic memory management</p>
+        <p>📌 Note: User is responsible for obtaining commercial license from Flux Dev if using image commercially under their license.</p>
     </div>
     """)