Sun, 28 Jun 2026 04:18 - Add enhance stage

README.md

@@ -8,10 +8,10 @@ sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: apache-2.0
-version: 0.1.2
+version: 0.1.3
 short_description: Powerful image editing - supports one or two input images.
 ---
 
-Pro Realism Edit Studio is a powerful image editor powered by [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer for 4-step inference. Upload one or two input images, write a prompt, get high-quality results.
+Pro Realism Edit Studio is a powerful image editor powered by [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer for 4-step inference. Upload one or two input images, write a prompt, get high-quality results. Optional Enhance modes add lazy Nomos ATD upscaling, masked skin cleanup, and final detail grain.
 
 Video generation is disabled unless an owned Gradio video Space is configured with `VIDEO_SPACE_ID`.

app.py

@@ -6,7 +6,7 @@ import spaces
 
 from accelerate import init_empty_weights
 from collections import OrderedDict
-from PIL import Image
+from PIL import Image, ImageEnhance, ImageFilter
 from diffusers.models import QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel
 from diffusers.models.model_loading_utils import load_model_dict_into_meta
 from huggingface_hub import hf_hub_download
@@ -22,7 +22,7 @@ from gradio_client import Client, handle_file
 import tempfile
 
 BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
-APP_VERSION = "0.1.2"
+APP_VERSION = "0.1.3"
 PHR00T_REPO_ID = os.environ.get("PHR00T_REPO_ID", "Phr00t/Qwen-Image-Edit-Rapid-AIO").strip()
 RAPID_TRANSFORMER_FILENAME = os.environ.get(
     "RAPID_TRANSFORMER_FILENAME",
@@ -30,6 +30,20 @@ RAPID_TRANSFORMER_FILENAME = os.environ.get(
 ).strip()
 PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."
 VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()
+UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "Phips/4xNomos8k_atd_jpg").strip()
+UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "4xNomos8k_atd_jpg.safetensors").strip()
+UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
+UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "32"))
+ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "1280"))
+ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.018"))
+
+ENHANCE_MODE_OFF = "Off"
+ENHANCE_MODE_UPSCALE = "Upscale"
+ENHANCE_MODE_CLEAN = "Clean"
+ENHANCE_MODE_MAX_DETAIL = "Max Detail"
+ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]
+
+_upscaler_model = None
 
 def turn_into_video(input_image, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
     if not VIDEO_SPACE_ID:
@@ -179,6 +193,141 @@ pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 
+
+def load_upscaler_model():
+    global _upscaler_model
+    if _upscaler_model is not None:
+        return _upscaler_model
+
+    try:
+        import spandrel
+        import spandrel_extra_arches
+    except ImportError as exc:
+        raise gr.Error("Enhance mode requires spandrel and spandrel_extra_arches to be installed.") from exc
+
+    spandrel_extra_arches.install()
+    model_path = hf_hub_download(repo_id=UPSCALER_MODEL_ID, filename=UPSCALER_MODEL_FILENAME)
+    model = spandrel.ModelLoader().load_from_file(model_path)
+    model.eval().to(device)
+    _upscaler_model = model
+    return _upscaler_model
+
+
+def image_to_tensor(image):
+    array = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor.to(device)
+
+
+def tensor_to_image(tensor):
+    array = tensor.squeeze(0).detach().float().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
+    return Image.fromarray((array * 255.0).round().astype(np.uint8), mode="RGB")
+
+
+def validate_enhance_input_size(image):
+    max_edge = max(image.size)
+    if max_edge > ENHANCE_MAX_INPUT_EDGE:
+        raise gr.Error(
+            f"Enhance mode accepts images up to {ENHANCE_MAX_INPUT_EDGE}px on the longest edge. "
+            f"Current image is {image.width}x{image.height}."
+        )
+
+
+def tile_upscale(image):
+    validate_enhance_input_size(image)
+    model = load_upscaler_model()
+    tensor = image_to_tensor(image)
+    _, _, height, width = tensor.shape
+    tile_size = max(64, min(UPSCALER_TILE_SIZE, height, width))
+    overlap = max(0, min(UPSCALER_TILE_OVERLAP, tile_size // 2))
+    step = max(1, tile_size - overlap)
+    y_positions = sorted(set(list(range(0, height, step)) + [max(0, height - tile_size)]))
+    x_positions = sorted(set(list(range(0, width, step)) + [max(0, width - tile_size)]))
+    output = None
+    weights = None
+
+    with torch.inference_mode():
+        for y in y_positions:
+            for x in x_positions:
+                y1 = min(y + tile_size, height)
+                x1 = min(x + tile_size, width)
+                tile = tensor[:, :, y:y1, x:x1]
+                upscaled_tile = model(tile).clamp(0, 1)
+                scale_y = upscaled_tile.shape[-2] // tile.shape[-2]
+                scale_x = upscaled_tile.shape[-1] // tile.shape[-1]
+                if output is None:
+                    output = torch.zeros(
+                        (1, 3, height * scale_y, width * scale_x),
+                        dtype=upscaled_tile.dtype,
+                        device=upscaled_tile.device,
+                    )
+                    weights = torch.zeros_like(output)
+                oy0, oy1 = y * scale_y, y1 * scale_y
+                ox0, ox1 = x * scale_x, x1 * scale_x
+                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile
+                weights[:, :, oy0:oy1, ox0:ox1] += 1
+
+    output = output / weights.clamp_min(1)
+    return tensor_to_image(output)
+
+
+def skin_repair_mask(image):
+    ycbcr = np.asarray(image.convert("YCbCr"))
+    cb = ycbcr[:, :, 1]
+    cr = ycbcr[:, :, 2]
+    mask = (
+        (cr >= 135)
+        & (cr <= 180)
+        & (cb >= 75)
+        & (cb <= 135)
+    ).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask, mode="L")
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.2))
+
+
+def repair_skin_texture(image):
+    base = image.convert("RGB")
+    mask = skin_repair_mask(base)
+    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.35))
+    blended = Image.composite(repaired, base, mask)
+    return ImageEnhance.Sharpness(blended).enhance(1.08)
+
+
+def add_film_grain(image, seed):
+    base = image.convert("RGB")
+    array = np.asarray(base).astype(np.float32)
+    rng = np.random.default_rng(seed)
+    grain = rng.normal(0.0, 255.0 * ENHANCE_GRAIN_STRENGTH, size=(array.shape[0], array.shape[1], 1))
+    array = np.clip(array + grain, 0, 255)
+    return Image.fromarray(array.astype(np.uint8), mode="RGB")
+
+
+def apply_enhancement(image, enhance_mode, seed=0, progress=None):
+    mode = enhance_mode or ENHANCE_MODE_OFF
+    if mode not in ENHANCE_MODE_CHOICES:
+        raise gr.Error(f"Unknown enhance mode: {mode}")
+    if mode == ENHANCE_MODE_OFF:
+        return image
+
+    enhanced = image.convert("RGB")
+    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL):
+        if progress:
+            progress(0.76, desc="Repairing skin texture...")
+        enhanced = repair_skin_texture(enhanced)
+
+    if mode in (ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL):
+        if progress:
+            progress(0.82, desc="Upscaling image...")
+        enhanced = tile_upscale(enhanced)
+
+    if mode == ENHANCE_MODE_MAX_DETAIL:
+        if progress:
+            progress(0.92, desc="Adding final grain and sharpness...")
+        enhanced = ImageEnhance.Sharpness(enhanced).enhance(1.12)
+        enhanced = add_film_grain(enhanced, seed)
+
+    return enhanced
+
 def use_output_as_input(output_images):
     """Move the first output image into the Image 1 slot."""
     if not output_images:
@@ -200,6 +349,7 @@ def infer(
     num_inference_steps=4,
     height=None,
     width=None,
+    enhance_mode=ENHANCE_MODE_OFF,
     num_images_per_prompt=1,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -249,6 +399,12 @@ def infer(
         num_images_per_prompt=num_images_per_prompt,
     ).images
 
+    if enhance_mode != ENHANCE_MODE_OFF:
+        images_pil = [
+            apply_enhancement(img, enhance_mode, seed=seed + idx, progress=progress)
+            for idx, img in enumerate(images_pil)
+        ]
+
     # Save images to temporary files for proper serving
     output_paths = []
     os.makedirs("outputs", exist_ok=True)
@@ -306,6 +462,12 @@ with gr.Blocks(css=css) as demo:
                     show_label=True,
                     placeholder="Enter your prompt here...",
             )
+                enhance_mode = gr.Radio(
+                    label="Enhance / Amelioration",
+                    choices=ENHANCE_MODE_CHOICES,
+                    value=ENHANCE_MODE_OFF,
+                    interactive=True,
+                )
                 run_button = gr.Button("Edit!", variant="primary")
                 
                 with gr.Accordion("Advanced Settings", open=False):
@@ -392,6 +554,7 @@ with gr.Blocks(css=css) as demo:
             num_inference_steps,
             height,
             width,
+            enhance_mode,
         ],
         outputs=[result, seed, use_output_btn, turn_video_btn],
 

requirements.txt

@@ -9,3 +9,5 @@ kernels==0.11.0
 torchvision
 peft
 torchao==0.11.0
+spandrel
+spandrel_extra_arches

tests/test_enhance_stage.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+import unittest
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def read_text(relative_path: str) -> str:
+    return (ROOT / relative_path).read_text(encoding="utf-8")
+
+
+class EnhanceStageContractTest(unittest.TestCase):
+    def test_enhance_modes_match_product_contract(self):
+        app = read_text("app.py")
+
+        self.assertIn('ENHANCE_MODE_OFF = "Off"', app)
+        self.assertIn('ENHANCE_MODE_UPSCALE = "Upscale"', app)
+        self.assertIn('ENHANCE_MODE_CLEAN = "Clean"', app)
+        self.assertIn('ENHANCE_MODE_MAX_DETAIL = "Max Detail"', app)
+        self.assertIn(
+            "ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]",
+            app,
+        )
+        self.assertIn("enhance_mode = gr.Radio", app)
+        self.assertIn("apply_enhancement", app)
+
+    def test_upscaler_is_lazy_and_configurable(self):
+        app = read_text("app.py")
+
+        self.assertIn("load_upscaler_model", app)
+        self.assertIn("UPSCALER_MODEL_ID", app)
+        self.assertIn("UPSCALER_MODEL_FILENAME", app)
+        self.assertIn("hf_hub_download(", app)
+        self.assertIn("spandrel", app)
+        self.assertIn("tile_upscale", app)
+
+    def test_upscaler_dependencies_are_declared(self):
+        requirements = read_text("requirements.txt")
+
+        self.assertIn("spandrel", requirements)
+        self.assertIn("spandrel_extra_arches", requirements)
+
+
+if __name__ == "__main__":
+    unittest.main()