Update app.py

app.py

@@ -1,5 +1,26 @@
 import os
 import gc
+import subprocess
+
+# ── DLSS 5 patch for Flux2KleinKV ──
+def _apply_dlss_patch():
+    try:
+        import diffusers
+        site_packages = os.path.dirname(diffusers.__file__)
+        patch_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "flux2_klein_kv.patch")
+        if os.path.exists(patch_file):
+            result = subprocess.run(
+                ["patch", "-p2", "--forward", "--batch"],
+                cwd=os.path.dirname(site_packages),
+                stdin=open(patch_file),
+                capture_output=True, text=True,
+            )
+            print("DLSS patch:", result.stdout or "already applied")
+    except Exception as e:
+        print(f"DLSS patch warning: {e}")
+
+_apply_dlss_patch()
+
 import gradio as gr
 import numpy as np
 import spaces
@@ -21,6 +42,7 @@ print("torch.__version__ =", torch.__version__)
 print("Using device:", device)
 
 from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import ImageDraw, ImageFont
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
@@ -28,7 +50,7 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 dtype = torch.bfloat16
 
 pipe = QwenImageEditPlusPipeline.from_pretrained(
-    "Qwen/Qwen-Image-Edit-2511",
+    "FireRedTeam/FireRed-Image-Edit-1.1",
     transformer=QwenImageTransformer2DModel.from_pretrained(
         "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
         torch_dtype=dtype,
@@ -43,6 +65,92 @@ try:
 except Exception as e:
     print(f"Warning: Could not set FA3 processor: {e}")
 
+# ── DLSS 5 model (Flux2KleinKV) ──
+_dlss_pipe = None
+def _get_dlss_pipe():
+    global _dlss_pipe
+    if _dlss_pipe is None:
+        try:
+            from diffusers.pipelines.flux2.pipeline_flux2_klein_kv import Flux2KleinKVPipeline
+            HF_TOKEN = os.environ.get("HF_TOKEN")
+            print("Loading DLSS 5 (FLUX.2-klein-9b-kv)...")
+            _dlss_pipe = Flux2KleinKVPipeline.from_pretrained(
+                "black-forest-labs/FLUX.2-klein-9b-kv",
+                torch_dtype=torch.bfloat16,
+                token=HF_TOKEN,
+            ).to(device)
+            print("DLSS 5 model loaded.")
+        except Exception as e:
+            print(f"DLSS model load failed: {e}")
+            _dlss_pipe = None
+    return _dlss_pipe
+
+# ── DLSS font helper ──
+_DLSS_FONT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Inter-Bold.ttf")
+def _get_dlss_font(size):
+    try:
+        return ImageFont.truetype(_DLSS_FONT_PATH, size)
+    except Exception:
+        return ImageFont.load_default(size=size)
+
+def _dlss_comparison(original: "Image.Image", enhanced: "Image.Image") -> "Image.Image":
+    """Side-by-side comparison image with DLSS 5 Off/On labels."""
+    w, h = original.size
+    enhanced = enhanced.resize((w, h), LANCZOS)
+    canvas = Image.new("RGB", (w * 2, h))
+    canvas.paste(original, (0, 0))
+    canvas.paste(enhanced, (w, 0))
+    overlay = Image.new("RGBA", (w * 2, h), (0, 0, 0, 0))
+    draw = ImageDraw.Draw(overlay)
+    font_size = max(16, int(h * 0.06))
+    font = _get_dlss_font(font_size)
+    pad_x = int(font_size * 1.0)
+    pad_y = int(font_size * 0.55)
+    margin_bottom = int(h * 0.06)
+    for label, cx, dark, green_bar in [
+        ("Before Edit", w // 2, True, False),
+        ("DLSS 5 On",   w + w // 2, False, True),
+    ]:
+        bbox = font.getbbox(label)
+        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
+        lw, lh = tw + 2 * pad_x, th + 2 * pad_y
+        gh = max(4, int(lh * 0.13)) if green_bar else 0
+        x = cx - lw // 2
+        y = h - margin_bottom - lh - gh
+        if dark:
+            draw.rectangle([x, y, x+lw, y+lh], fill=(10,10,10,225), outline=(75,75,75,255), width=1)
+            draw.text((x + lw//2, y + lh//2), label, fill=(255,255,255,255), font=font, anchor="mm")
+        else:
+            draw.rectangle([x, y, x+lw, y+lh], fill=(255,255,255,255), outline=(190,190,190,255), width=1)
+            draw.text((x + lw//2, y + lh//2), label, fill=(0,0,0,255), font=font, anchor="mm")
+            draw.rectangle([x, y+lh, x+lw, y+lh+gh], fill=(118,185,0,255))
+    canvas = Image.alpha_composite(canvas.convert("RGBA"), overlay)
+    return canvas.convert("RGB")
+
+def _run_dlss(image: "Image.Image", seed: int) -> "Image.Image":
+    """Run DLSS 5 upscale/enhance on an image. Returns enhanced image (same size)."""
+    dlss = _get_dlss_pipe()
+    if dlss is None:
+        return image  # graceful fallback: return original if model didn't load
+    iw, ih = image.size
+    ar = iw / ih
+    if ar >= 1:
+        width = 1024; height = round(1024 / ar / 8) * 8
+    else:
+        height = 1024; width = round(1024 * ar / 8) * 8
+    width = max(256, min(1024, width))
+    height = max(256, min(1024, height))
+    generator = torch.Generator(device=device).manual_seed(seed)
+    result = dlss(
+        prompt="make it more realistic",
+        image=[image],
+        height=height,
+        width=width,
+        num_inference_steps=4,
+        generator=generator,
+    ).images[0]
+    return result
+
 EXAMPLES_CONFIG = [
     {
         "images": ["examples/1.jpg"],
@@ -195,11 +303,21 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
     negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
     width, height = update_dimensions_on_upload(pil_images[0])
     try:
-        result_image = pipe(
+        edited = pipe(
             image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
             height=height, width=width, num_inference_steps=steps,
             generator=generator, true_cfg_scale=guidance_scale,
         ).images[0]
+
+        # Auto-DLSS: upscale/enhance the edited result
+        dlss_seed = seed + 1 if not randomize_seed else seed
+        try:
+            dlss_enhanced = _run_dlss(edited, dlss_seed)
+            result_image = _dlss_comparison(edited, dlss_enhanced)
+        except Exception as dlss_err:
+            print(f"DLSS step skipped: {dlss_err}")
+            result_image = edited  # fallback to plain edited image
+
         return result_image, seed
     except Exception as e:
         raise e