Update app.py

app.py

@@ -5,9 +5,15 @@ import numpy as np
 import spaces
 import torch
 import random
+import cv2
 from PIL import Image
 from typing import Iterable
 
+# ── requirements.txt should include: ─────────────────────────────────────────
+#   basicsr
+#   realesrgan
+# ─────────────────────────────────────────────────────────────────────────────
+
 from diffusers import Flux2KleinPipeline
 from huggingface_hub import hf_hub_download
 
@@ -225,7 +231,7 @@ def on_gallery_change(images):
         img = Image.open(path if isinstance(path, str) else path.name)
         orig_w, orig_h = img.size
         base_w, base_h = compute_base_dimensions(img)
-        size_text = f"Input: **{orig_w} × {orig_h}** px → Base output (1×): **{base_w} × {base_h}** px"
+        size_text = f"Input: **{orig_w} × {orig_h}** px → Output (pre-upscale): **{base_w} × {base_h}** px"
         return count_info, size_text
     except Exception as e:
         return count_info, f"*Could not read dimensions: {e}*"
@@ -293,6 +299,47 @@ def update_weight_sliders(selected_titles):
 
 # ── Inference ─────────────────────────────────────────────────────────────────
 
+def apply_realesrgan(image: Image.Image, scale: int) -> Image.Image:
+    """
+    Upscale a PIL image using Real-ESRGAN.
+    scale=2 uses RealESRGAN_x2plus; scale=4 uses RealESRGAN_x4plus.
+    Imports are lazy so the packages are only required when upscaling is used.
+    """
+    try:
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+    except ImportError:
+        raise gr.Error(
+            "Real-ESRGAN is not installed. Add 'basicsr' and 'realesrgan' to requirements.txt."
+        )
+
+    if scale == 2:
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
+        sr_scale = 2
+    else:  # 4×
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+        sr_scale = 4
+
+    upsampler = RealESRGANer(
+        scale=sr_scale,
+        model_path=model_url,
+        model=model,
+        tile=512,          # tile to keep VRAM usage manageable
+        tile_pad=10,
+        pre_pad=0,
+        half=True,         # bfloat16/fp16 — matches the rest of the pipeline
+        device=device,
+    )
+
+    img_np = np.array(image)
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    output_bgr, _ = upsampler.enhance(img_bgr, outscale=sr_scale)
+    output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(output_rgb)
+
+
 @spaces.GPU
 def infer(
     input_images,
@@ -303,7 +350,7 @@ def infer(
     randomize_seed,
     guidance_scale,
     steps,
-    output_scale,
+    upscale_factor,
     *slider_values,
     progress=gr.Progress(track_tqdm=True)
 ):
@@ -365,11 +412,9 @@ def infer(
     else:
         full_prompt = user_prompt or lora_extra
 
-    # Compute output dimensions: base size × output_scale, aligned to 16px grid
-    base_w, base_h = compute_base_dimensions(pil_images[0])
-    width  = max(16, (int(base_w * output_scale) // 16) * 16)
-    height = max(16, (int(base_h * output_scale) // 16) * 16)
-    print(f"Output size: {width} × {height} px (scale {output_scale}×)")
+    # Always generate at base 1024-capped resolution
+    width, height = compute_base_dimensions(pil_images[0])
+    print(f"Generating at: {width} × {height} px")
 
     processed_images = [img.resize((width, height), Image.LANCZOS).convert("RGB") for img in pil_images]
     image_input = processed_images if len(processed_images) > 1 else processed_images[0]
@@ -384,12 +429,23 @@ def infer(
             num_inference_steps=steps,
             generator=torch.Generator(device=device).manual_seed(seed),
         ).images[0]
-        return image, seed
     except Exception as e:
         raise gr.Error(f"Inference failed: {e}")
-    finally:
-        gc.collect()
-        torch.cuda.empty_cache()
+
+    # ── Post-generation upscaling ─────────────────────────────────────────────
+    if upscale_factor and upscale_factor != "None":
+        scale_int = int(upscale_factor[0])  # "2×" → 2, "4×" → 4
+        print(f"Upscaling {scale_int}× with Real-ESRGAN…")
+        try:
+            image = apply_realesrgan(image, scale_int)
+            print(f"Upscaled to: {image.width} × {image.height} px")
+        except Exception as e:
+            gr.Warning(f"Upscaling failed, returning 1024px result: {e}")
+    # ─────────────────────────────────────────────────────────────────────────
+
+    gc.collect()
+    torch.cuda.empty_cache()
+    return image, str(seed)
 
 # ── Dynamic LoRA loader ───────────────────────────────────────────────────────
 
@@ -482,16 +538,23 @@ with gr.Blocks(css=css, theme=orange_red_theme) as demo:
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
-                    output_scale = gr.Slider(
-                        label="Output Scale",
-                        info="Multiplier applied to the auto-detected base size shown above. 1× keeps the default behaviour.",
-                        minimum=0.25, maximum=4.0, step=0.25, value=1.0,
+                    upscale_factor = gr.Radio(
+                        label="Upscale (Real-ESRGAN)",
+                        info="Applied after generation. 2× ≈ 2048px, 4× ≈ 4096px on the longest side.",
+                        choices=["None", "2×", "4×"],
+                        value="None",
                     )
 
             # ── Right column ─────────────────────────────────────────────────
             with gr.Column(scale=1):
                 output_image = gr.Image(label="Output", interactive=False, format="png", height=320)
-                used_seed = gr.Textbox(label="Used Seed", interactive=False, visible=False)
+                used_seed = gr.Textbox(
+                    label="🌱 Seed used",
+                    info="Re-enter this seed in Advanced Settings to reproduce the result exactly.",
+                    interactive=False,
+                    visible=True,
+                    max_lines=1,
+                )
 
         # ── LoRA selector ─────────────────────────────────────────────────────
         gr.Markdown("### 🎨 Select LoRA(s)")
@@ -537,7 +600,7 @@ with gr.Blocks(css=css, theme=orange_red_theme) as demo:
             ],
             inputs=[input_images, prompt, lora_selector, seed, randomize_seed, guidance_scale, steps],
             outputs=[output_image, used_seed],
-            fn=lambda imgs, p, sel, sd, rs, gs, st: infer(imgs, p, "", sel, sd, rs, gs, st, 1.0, *([1.0] * MAX_LORA_SLOTS)),
+            fn=lambda imgs, p, sel, sd, rs, gs, st: infer(imgs, p, "", sel, sd, rs, gs, st, "None", *([1.0] * MAX_LORA_SLOTS)),
             cache_examples=False,
             label="Examples",
         )
@@ -566,7 +629,7 @@ with gr.Blocks(css=css, theme=orange_red_theme) as demo:
 
     run_button.click(
         fn=infer,
-        inputs=[input_images, prompt, lora_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, output_scale] + weight_sliders,
+        inputs=[input_images, prompt, lora_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor] + weight_sliders,
         outputs=[output_image, used_seed],
     )