Upload app.py

app.py

@@ -68,7 +68,6 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
 import Age_classifier
 
-from esrgan import load_esrgan, upscale_to_pixels
 from security import (
     FloodGuard,
     check_admin_password,
@@ -120,10 +119,40 @@ if torch.cuda.is_available():
 MODEL_QWEN = "qwen"
 MODEL_FLUX = "flux"
 
-ASPECT_RATIOS = ["Default", "1:1", "3:2", "2:3", "9:16", "16:9", "3:4", "4:3", "21:9", "9:21"]
-RATIO_DEFAULT = "Default"
-QWEN_TARGET_PIXELS = 2073600
-UPSCALE_PIXELS = {"Basic": 3686400, "Medium": 5760000, "High": 8294400}
+PRESET_RESOLUTIONS = [
+    ("720×1280 (9:16)", 720, 1280),
+    ("864×1536 (9:16)", 864, 1536),
+    ("1080×1920 (9:16)", 1080, 1920),
+    ("1440×2560 (9:16)", 1440, 2560),
+    ("1600×2560 (10:16)", 1600, 2560),
+    ("1536×2728 (9:16)", 1536, 2728),
+    ("1280×720 (16:9)", 1280, 720),
+    ("1536×864 (16:9)", 1536, 864),
+    ("1920×1080 (16:9)", 1920, 1080),
+    ("2560×1440 (16:9)", 2560, 1440),
+    ("2560×1600 (16:10)", 2560, 1600),
+    ("2728×1536 (16:9)", 2728, 1536),
+    ("768×1024 (3:4)", 768, 1024),
+    ("1080×1440 (3:4)", 1080, 1440),
+    ("1536×2048 (3:4)", 1536, 2048),
+    ("1024×1536 (2:3)", 1024, 1536),
+    ("1440×1800 (4:5)", 1440, 1800),
+    ("1080×2400 (9:20)", 1080, 2400),
+    ("1024×768 (4:3)", 1024, 768),
+    ("1440×1080 (4:3)", 1440, 1080),
+    ("2048×1536 (4:3)", 2048, 1536),
+    ("1536×1024 (3:2)", 1536, 1024),
+    ("1800×1440 (5:4)", 1800, 1440),
+    ("2560×1080 (21:9)", 2560, 1080),
+]
+PRESET_RESOLUTION_MAP = {label: (w, h) for label, w, h in PRESET_RESOLUTIONS}
+PRESET_TALL = [label for label, w, h in PRESET_RESOLUTIONS if w < h]
+PRESET_WIDE = [label for label, w, h in PRESET_RESOLUTIONS if w > h]
+PRESET_ALL = PRESET_TALL + PRESET_WIDE
+PRESET_DEFAULT = "720×1280 (9:16)"
+AUTO_SCALES = ["1x", "2x", "3x", "4x"]
+AUTO_BASE_PIXELS = 1024 * 1024
+AUTO_MAX_PIXELS = 4194304
 
 from diffusers.utils import logging as diffusers_logging
 
@@ -219,7 +248,7 @@ def _enable_vae_tiling(p, tile=None):
         print("[vae] tiling applied:", {k: getattr(vae, k, None) for k in ("tile_sample_min_height", "tile_sample_min_width", "tile_sample_stride_height", "tile_sample_stride_width")})
 
 
-QWEN_VAE_TILE = int(os.environ.get("QWEN_VAE_TILE", "0"))
+QWEN_VAE_TILE = int(os.environ.get("QWEN_VAE_TILE", "1600"))
 
 
 pipe = QwenImageEditPlusPipeline.from_pretrained(
@@ -256,20 +285,16 @@ if getattr(pipe, "text_encoder", None) is not None:
     except Exception as e:
         print(f"[qwen] text_encoder quantization FAILED, staying bf16 (~+8GB VRAM): {e}")
 
-if QWEN_VAE_TILE > 0:
-    _enable_vae_tiling(pipe, QWEN_VAE_TILE)
-else:
-    print("[vae] Qwen VAE tiling OFF (full-frame decode)")
+_enable_vae_tiling(pipe, QWEN_VAE_TILE)
 
 pipe.transformer.__class__ = QwenImageTransformer2DModel
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
 FLUX_REPO = "black-forest-labs/FLUX.2-klein-9B"
-FLUX_REFINE_PROMPT = "An ultra-high-resolution masterpiece with exceptional sharpness and detail, featuring only a subtle, barely perceptible amount of noise on the skin’s surface. The image remains entirely unzoomed and unaltered, preserving its original composition and proportions, with lighting and shadows meticulously controlled to remain perfectly faithful to the original."
+FLUX_REFINE_PROMPT = "An ultra-high-resolution masterpiece with exceptional sharpness and detail, featuring only a subtle, barely perceptible amount of noise on the skin’s surface. The image remains entirely unzoomed and unaltered, preserving its original composition and proportions, with lighting and shadows meticulously controlled to remain perfectly faithful to the original"
 FLUX_REFINE_GUIDANCE_SCALE = 1
 FLUX_NUM_INFERENCE_STEPS = 4
-FLUX_SEED = 42
-FLUX_MAX_PIXELS = 2073600
+FLUX_MAX_PIXELS = 2048 * 2048
 
 flux_pipe = Flux2KleinPipeline.from_pretrained(FLUX_REPO, torch_dtype=dtype)
 
@@ -288,33 +313,6 @@ try:
 except Exception as e:
     print(f"[flux] quantization FAILED, continuing in bf16 (roughly 2x the VRAM/transfer time): {e}")
 
-if hasattr(flux_pipe, "vae") and flux_pipe.vae is not None:
-    try:
-        flux_pipe.vae.enable_slicing()
-    except Exception:
-        pass
-
-_flux_ip = getattr(flux_pipe, "image_processor", None)
-if _flux_ip is not None and hasattr(_flux_ip, "_resize_to_target_area"):
-    _orig_resize_to_target_area = _flux_ip._resize_to_target_area
-
-    def _flux_resize_with_higher_cap(image, target_area):
-        if target_area == 1024 * 1024:
-            w, h = image.size
-            if w * h <= FLUX_MAX_PIXELS:
-                return image
-            return _orig_resize_to_target_area(image, FLUX_MAX_PIXELS)
-        return _orig_resize_to_target_area(image, target_area)
-
-    _flux_ip._resize_to_target_area = _flux_resize_with_higher_cap
-    print(f"[flux] reference-image cap raised: no downscale up to {FLUX_MAX_PIXELS} px")
-else:
-    print("[flux] WARNING: image_processor._resize_to_target_area not found; "
-          "FLUX will use its built-in ~1MP reference-image downscale")
-
-esrgan_model = load_esrgan(device)
-print("[esrgan] RealESRGAN_x2plus loaded")
-
 # FLUX VAE tiling is off by default because tiled encode/decode shifted skin colour.
 if os.environ.get("FLUX_VAE_TILING", "0") == "1":
     _enable_vae_tiling(flux_pipe)
@@ -323,6 +321,12 @@ gc.collect()
 
 
 MAX_SEED = 2**31 - 1
+LIMIT_STEPS_MAX = 20
+LIMIT_GUIDANCE_MIN = 1.0
+LIMIT_GUIDANCE_MAX = 2.0
+LIMIT_SIDE_MAX = 2728
+LIMIT_PIXELS_MAX = 2048 * 2048
+LIMIT_INPUT_PIXELS = 2560 * 2560
 _VALID_FORMATS = ("PNG", "JPG", "WebP")
 
 GPU_DURATION_MIN = 10
@@ -342,10 +346,6 @@ def _to_int(value, lo, hi, default):
     return max(lo, min(hi, v))
 
 
-LIMIT_GUIDANCE_MIN = 1.0
-LIMIT_GUIDANCE_MAX = 2.0
-
-
 def _to_float(value, lo, hi, default):
     try:
         v = float(value)
@@ -356,38 +356,42 @@ def _to_float(value, lo, hi, default):
     return max(lo, min(hi, v))
 
 
-def _round_size(width, height, multiple=8):
-    return max(multiple, (int(width) // multiple) * multiple), max(multiple, (int(height) // multiple) * multiple)
-
+def _clamp_size(width, height, multiple=8, lo=64):
+    w = _to_int(width, lo, LIMIT_SIDE_MAX, lo)
+    h = _to_int(height, lo, LIMIT_SIDE_MAX, lo)
+    if w * h > LIMIT_PIXELS_MAX:
+        scale = (LIMIT_PIXELS_MAX / (w * h)) ** 0.5
+        w, h = int(w * scale), int(h * scale)
+    return max(lo, (w // multiple) * multiple), max(lo, (h // multiple) * multiple)
 
-def _ratio_size(ratio, pixels=QWEN_TARGET_PIXELS, multiple=16):
-    width = (pixels * ratio) ** 0.5
-    height = width / ratio
-    return (
-        max(multiple, int(width) // multiple * multiple),
-        max(multiple, int(height) // multiple * multiple),
-    )
 
-
-def _image_ratio(image_1, image_2):
+def _auto_size(scale_choice, image_1, image_2):
     src = image_1 if image_1 is not None else image_2
     if src is None:
-        return 1.0
+        return None, None
     try:
         with Image.open(src if isinstance(src, str) else src.name) as im:
-            return im.width / im.height
+            width, height = im.size
     except Exception:
-        return 1.0
+        return None, None
+    factor = _to_int(str(scale_choice).rstrip("x"), 1, 4, 1)
+    pixels = min(AUTO_BASE_PIXELS * factor, AUTO_MAX_PIXELS)
+    ratio = width / height
+    out_w = (pixels * ratio) ** 0.5
+    return _clamp_size(int(out_w), int(out_w / ratio), multiple=16)
 
 
-def _resolve_size(ratio_label, image_1, image_2):
+def _resolve_size(mode, preset, custom_width, custom_height, auto_scale, image_1, image_2):
     try:
-        if ratio_label in ASPECT_RATIOS and ratio_label != RATIO_DEFAULT:
-            w, h = ratio_label.split(":")
-            return _ratio_size(int(w) / int(h))
-    except (TypeError, ValueError, ZeroDivisionError):
+        if mode == "Preset" and preset in PRESET_RESOLUTION_MAP:
+            return _clamp_size(*PRESET_RESOLUTION_MAP[preset])
+        if mode == "Custom":
+            return _clamp_size(int(custom_width), int(custom_height))
+        if mode == "Auto":
+            return _auto_size(auto_scale, image_1, image_2)
+    except (TypeError, ValueError, OverflowError):
         pass
-    return _ratio_size(_image_ratio(image_1, image_2))
+    return None, None
 
 
 def _fit_size(width, height, max_pixels=FLUX_MAX_PIXELS, multiple=16):
@@ -403,7 +407,12 @@ def _fit_size(width, height, max_pixels=FLUX_MAX_PIXELS, multiple=16):
 def _load_image(img):
     if not isinstance(img, Image.Image):
         img = Image.open(img if isinstance(img, str) else img.name)
-    return img.convert("RGB")
+    img = img.convert("RGB")
+    w, h = img.size
+    if w * h > LIMIT_INPUT_PIXELS:
+        scale = (LIMIT_INPUT_PIXELS / (w * h)) ** 0.5
+        img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
+    return img
 
 
 def _sanitize_common(seed, gpu_duration, output_format):
@@ -447,6 +456,7 @@ def _generate_core(
     num_inference_steps,
     height,
     width,
+    output_format,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -535,7 +545,6 @@ def _finalize(res, image_1, image_2, prompt, output_format):
         else:
             img.save(output_path)
         output_paths.append(output_path)
-        print(f"[out] resolution={img.width}x{img.height}")
         log_generation(
             path=output_path,
             prompt=prompt,
@@ -550,49 +559,30 @@ def _finalize(res, image_1, image_2, prompt, output_format):
             inputs=input_paths,
         )
     has_output = bool(output_paths)
-    return output_paths, seed, gr.update(visible=has_output)
+    return output_paths, seed, gr.update(visible=has_output), gr.update(visible=has_output)
 
 
-def _make_qwen_flux_gpu_variant(duration_seconds, size):
-    def _qwen_then_flux_gpu_call(qwen_kwargs, scale_up):
-        _free_vram_cache()
-        q = _generate_core(MODEL_QWEN, **qwen_kwargs)
-        if not q["images"]:
-            return q, None
-        _free_vram_cache()
-        img = q["images"][0]
-        w, h = _round_size(*img.size, multiple=16)
-        f = _generate_core(
-            MODEL_FLUX, image_1=img, image_2=None, prompt=FLUX_REFINE_PROMPT,
-            seed=FLUX_SEED, randomize_seed=False,
-            true_guidance_scale=FLUX_REFINE_GUIDANCE_SCALE,
-            num_inference_steps=FLUX_NUM_INFERENCE_STEPS,
-            height=h, width=w,
-            progress=qwen_kwargs["progress"],
-        )
-        if scale_up and f["images"]:
-            _free_vram_cache()
-            f = {**f, "images": [upscale_to_pixels(esrgan_model, f["images"][0], scale_up)]}
-        return q, f
+def _make_gpu_variant(duration_seconds, size):
+    def _infer_gpu_call(*args, **kwargs):
+        return _generate_core(*args, **kwargs)
 
-    _qwen_then_flux_gpu_call.__name__ = f"_qwen_then_flux_gpu_{size}_{duration_seconds}s"
-    _qwen_then_flux_gpu_call.__qualname__ = _qwen_then_flux_gpu_call.__name__
-    return spaces.GPU(duration=duration_seconds, size=size)(_qwen_then_flux_gpu_call)
+    _infer_gpu_call.__name__ = f"_infer_impl_gpu_{size}_{duration_seconds}s"
+    _infer_gpu_call.__qualname__ = _infer_gpu_call.__name__
+    return spaces.GPU(duration=duration_seconds, size=size)(_infer_gpu_call)
 
 
-_qwen_then_flux_variants = {(s, d): _make_qwen_flux_gpu_variant(d, s) for s in GPU_SIZE_OPTIONS for d in _GPU_DURATION_STEPS}
+_gpu_variants = {(s, d): _make_gpu_variant(d, s) for s in GPU_SIZE_OPTIONS for d in _GPU_DURATION_STEPS}
 
 
-def _pick_qwen_flux_variant(gpu_duration, gpu_size=GPU_SIZE_DEFAULT):
+def _pick_gpu_variant(gpu_duration, gpu_size=GPU_SIZE_DEFAULT):
     requested = _to_int(gpu_duration, GPU_DURATION_MIN, GPU_DURATION_MAX, GPU_DURATION_DEFAULT)
     chosen = next((d for d in _GPU_DURATION_STEPS if d >= requested), _GPU_DURATION_STEPS[-1])
     size = gpu_size if gpu_size in GPU_SIZE_OPTIONS else GPU_SIZE_DEFAULT
-    gpu_fn = _qwen_then_flux_variants[(size, chosen)]
+    gpu_fn = _gpu_variants[(size, chosen)]
 
-    def _call(image_1, image_2, prompt, output_format, qwen_kwargs, scale_up):
-        q, f = gpu_fn(qwen_kwargs, scale_up)
-        final = f if (f is not None and f["images"]) else q
-        return _finalize({**q, "images": final["images"]}, image_1, image_2, prompt, output_format)
+    def _call(model_choice, image_1, image_2, prompt, *args, output_format="PNG", **kwargs):
+        res = gpu_fn(model_choice, image_1, image_2, prompt, *args, output_format=output_format, **kwargs)
+        return _finalize(res, image_1, image_2, prompt, output_format)
 
     return _call
 
@@ -691,7 +681,7 @@ def _run_limited(request, seed, run):
     ok, key, msg = rate_acquire(request)
     if not ok:
         show_rate_message(msg)
-        return gr.update(), seed, gr.update()
+        return gr.update(), seed, gr.update(), gr.update()
     refund = False
     try:
         return run()
@@ -702,7 +692,7 @@ def _run_limited(request, seed, run):
         rate_release(key, refund=refund)
 
 
-def run_qwen_and_flux(
+def infer_qwen_only(
     image_1,
     image_2,
     prompt,
@@ -712,10 +702,11 @@ def run_qwen_and_flux(
     num_inference_steps,
     gpu_duration,
     gpu_size,
-    aspect_ratio,
-    scale_basic,
-    scale_medium,
-    scale_high,
+    resolution_mode,
+    resolution_preset,
+    custom_width,
+    custom_height,
+    auto_scale,
     output_format,
     age_confirmed=False,
     progress=gr.Progress(track_tqdm=True),
@@ -725,34 +716,77 @@ def run_qwen_and_flux(
         raise gr.Error("Please tick the box to confirm you are 18 or older.")
 
     if not (isinstance(prompt, str) and prompt.strip()):
-        return gr.update(), seed, gr.update()
+        return gr.update(), seed, gr.update(), gr.update()
 
     def run():
         clean_seed, duration, fmt = _sanitize_common(seed, gpu_duration, output_format)
-        width, height = _resolve_size(aspect_ratio, image_1, image_2)
+        width, height = _resolve_size(resolution_mode, resolution_preset, custom_width, custom_height, auto_scale, image_1, image_2)
         clean_prompt = prompt if isinstance(prompt, str) else ""
-        qwen_kwargs = dict(
-            image_1=image_1, image_2=image_2, prompt=clean_prompt,
-            seed=clean_seed, randomize_seed=bool(randomize_seed),
-            true_guidance_scale=_to_float(true_guidance_scale, LIMIT_GUIDANCE_MIN, LIMIT_GUIDANCE_MAX, 1.0),
-            num_inference_steps=int(num_inference_steps),
-            height=height, width=width,
-            progress=progress,
-        )
-        upscale_pixels = next(
-            (UPSCALE_PIXELS[name] for name, on in (("High", scale_high), ("Medium", scale_medium), ("Basic", scale_basic)) if on),
-            0,
-        )
         try:
             _screen_inputs(request, clean_prompt, image_1, image_2)
-            return _screen_outputs(
-                _pick_qwen_flux_variant(duration, gpu_size)(image_1, image_2, clean_prompt, fmt, qwen_kwargs, upscale_pixels),
-                request,
-            )
+            return _screen_outputs(_pick_gpu_variant(duration, gpu_size)(
+                MODEL_QWEN,
+                image_1,
+                image_2,
+                clean_prompt,
+                seed=clean_seed,
+                randomize_seed=bool(randomize_seed),
+                true_guidance_scale=_to_float(true_guidance_scale, LIMIT_GUIDANCE_MIN, LIMIT_GUIDANCE_MAX, 1.0),
+                num_inference_steps=_to_int(num_inference_steps, 1, LIMIT_STEPS_MAX, 4),
+                height=height,
+                width=width,
+                output_format=fmt,
+                progress=progress,
+            ), request)
+        except gr.Error as exc:
+            if not _is_age_block(exc):
+                raise
+            return [_block_image()], seed, gr.update(visible=False), gr.update(visible=False)
+
+    return _run_limited(request, seed, run)
+
+
+def refine_with_flux(
+    output_images,
+    seed,
+    randomize_seed,
+    gpu_duration,
+    gpu_size,
+    output_format,
+    progress=gr.Progress(track_tqdm=True),
+    request: gr.Request = None,
+):
+    if not output_images:
+        return gr.update(), seed, gr.update(visible=False), gr.update(visible=False)
+
+    def run():
+        clean_seed, duration, fmt = _sanitize_common(seed, gpu_duration, output_format)
+        path = _first_path(output_images)
+        try:
+            with Image.open(path) as im:
+                width, height = _clamp_size(*im.size, multiple=16)
+        except Exception:
+            width = height = None
+        try:
+            _screen_inputs(request, "", path)
+            return _screen_outputs(_pick_gpu_variant(duration, gpu_size)(
+                MODEL_FLUX,
+                path,
+                None,
+                FLUX_REFINE_PROMPT,
+                seed=clean_seed,
+                randomize_seed=bool(randomize_seed),
+                true_guidance_scale=FLUX_REFINE_GUIDANCE_SCALE,
+                num_inference_steps=FLUX_NUM_INFERENCE_STEPS,
+                height=height,
+                width=width,
+                output_format=fmt,
+                progress=progress,
+            ), request)
         except gr.Error as exc:
             if not _is_age_block(exc):
                 raise
-            return [_block_image()], seed, gr.update(visible=False)
+            return [_block_image()], seed, gr.update(visible=False), gr.update()
 
     return _run_limited(request, seed, run)
 
@@ -764,6 +798,11 @@ def use_output_as_input(output_images):
 
 
 def _sync_result_followups(output_images):
+    visible = bool(output_images) and not _is_block_image(output_images)
+    return gr.update(visible=visible), gr.update(visible=visible)
+
+
+def _sync_flux_result_followups(output_images):
     return gr.update(visible=bool(output_images) and not _is_block_image(output_images))
 
 
@@ -781,6 +820,23 @@ css = """
 }
 #logo-title {
     text-align: center;
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+    gap: 14px;
+    margin-bottom: 10px;
+}
+#logo-title h1 {
+    margin: 0 !important;
+    padding: 0 !important;
+    line-height: 1.15 !important;
+}
+/* Italic + background-clip:text clips the slanted tail of the last glyph; extra right padding
+   gives it room, and the negative margin keeps the title visually centered. */
+#logo-title h1#ast-main-title,
+#logo-title h1#ast-nfaa-title {
+    padding-right: 0.3em !important;
+    margin-right: -0.3em !important;
 }
 
 @keyframes ast-title-flash {
@@ -802,11 +858,6 @@ css = """
 @media (max-width: 480px) {
     #ast-main-title { font-size: 5.8vw !important; }
 }
-#logo-title h1#ast-nfaa-title {
-    margin: 0 -0.3em 0 0 !important;
-    padding: 0 0.3em 0 0 !important;
-    line-height: 1.15 !important;
-}
 #ast-nfaa-title {
     font-style: italic;
     display: inline-block;
@@ -826,7 +877,7 @@ css = """
 #ast-notice {
     width: 100%;
     max-width: 560px;
-    margin: 11px auto 0 !important;
+    margin: 0 auto !important;
     padding: 1px 0 1px 14px !important;
     border-left: 2px solid transparent;
     border-image: linear-gradient(180deg, #FF3CAC, #7B2FF7, #00E5FF) 1;
@@ -875,57 +926,63 @@ css = """
     background-color: #FF1010 !important;
     border-color: #FF1010 !important;
 }
+/* Force a 2-column grid for the preset radio (Gradio's default flex-wrap stacks into one column on mobile). */
+.preset-grid > .wrap,
+.preset-grid .wrap {
+    display: grid !important;
+    grid-template-columns: 1fr 1fr !important;
+    gap: 8px !important;
+}
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-custom-height,
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-custom-width,
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-auto-scale,
+#adv-settings:has(#res-mode input[value="Custom"]:checked) #res-preset,
+#adv-settings:has(#res-mode input[value="Custom"]:checked) #res-auto-scale,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-preset,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-custom-height,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-custom-width {
+    display: none !important;
+}
+/* Seed is an output component: hide Gradio's progress/status overlay so it never covers the number box. */
 #seed-slider .wrap.generating,
 #seed-slider .wrap.pending,
 #seed-slider .progress-text,
 #seed-slider .progress-level {
     display: none !important;
 }
-"""
-
-_RATIO_ICON_MAX = 30
-_RATIO_ICON_BOX = 34
-_RATIO_ACCENT = "var(--checkbox-background-color-selected, var(--color-accent, var(--primary-500)))"
-
-
-def _ratio_css():
-    acc = _RATIO_ACCENT
-    rules = [
-        "#aspect-ratio .wrap { display: grid !important; grid-template-columns: repeat(5, minmax(0, 1fr)) !important; gap: 6px !important; }",
-        "#aspect-ratio label { display: flex !important; flex-direction: column !important; align-items: center !important; "
-        "justify-content: center !important; gap: 0 !important; width: 100% !important; min-width: 0 !important; "
-        "margin: 0 !important; padding: 8px 2px 6px !important; box-sizing: border-box !important; "
-        "border-radius: 12px !important; cursor: pointer; }",
-        "#aspect-ratio label input { display: none !important; }",
-        "#aspect-ratio label span { font-size: 13px !important; line-height: 1.2 !important; white-space: nowrap; margin: 0 !important; padding: 0 !important; text-indent: 0 !important; text-align: center !important; }",
-        "#aspect-ratio label::before { content: \"\"; display: block; flex: none; border-radius: 4px; background: rgba(128, 128, 128, 0.38); }",
-        f"#aspect-ratio label:has(input:checked), #aspect-ratio label.selected {{ border-color: color-mix(in srgb, {acc} 90%, transparent) !important; background: color-mix(in srgb, {acc} 14%, transparent) !important; }}",
-        f"#aspect-ratio label:has(input:checked)::before, #aspect-ratio label.selected::before {{ background: color-mix(in srgb, {acc} 70%, transparent); }}",
-        '#aspect-ratio label:has(input[value="Default"])::before { background: transparent !important; border: 1.5px dashed rgba(150, 150, 150, 0.8); box-sizing: border-box; }',
-        f'#aspect-ratio label:has(input[value="Default"]:checked)::before, #aspect-ratio label.selected:has(input[value="Default"])::before {{ border-color: color-mix(in srgb, {acc} 95%, transparent); }}',
-    ]
-    for label in ASPECT_RATIOS:
-        w, h = (4, 3) if label == RATIO_DEFAULT else (int(x) for x in label.split(":"))
-        k = _RATIO_ICON_MAX / max(w, h)
-        iw, ih = round(w * k), round(h * k)
-        pad = (_RATIO_ICON_BOX - ih) / 2
-        rules.append(
-            f'#aspect-ratio label:has(input[value="{label}"])::before '
-            f'{{ width: {iw}px; height: {ih}px; margin: {pad:g}px 0; }}'
-        )
-    return "\n".join(rules) + "\n"
-
-
-css += _ratio_css()
-css += f"""
-#output-format {{ background: transparent !important; border: none !important; box-shadow: none !important; padding: 0 !important; }}
-#output-format .wrap {{ display: flex !important; justify-content: center !important; gap: 1px !important; background: transparent !important; padding: 0 !important; margin: 0 !important; }}
-#output-format label {{ flex: 1 1 0 !important; min-width: 0 !important; justify-content: center !important; align-items: center !important; height: 38px !important; min-height: 0 !important; margin: 0 !important; padding: 0 8px !important; border: none !important; border-radius: 0 !important; box-shadow: none !important; background: var(--block-background-fill) !important; }}
-#output-format label:has(input:checked), #output-format label.selected, #lanczos-row .block:has(input:checked) {{ background-image: linear-gradient(color-mix(in srgb, {_RATIO_ACCENT} 14%, transparent), color-mix(in srgb, {_RATIO_ACCENT} 14%, transparent)) !important; }}
-.lanczos-title {{ padding: 4px 0 3px 6px !important; background: var(--block-background-fill) !important; }}
-"""
-
-css += """
+#flux-seed-slider .wrap.generating,
+#flux-seed-slider .wrap.pending,
+#flux-seed-slider .progress-text,
+#flux-seed-slider .progress-level {
+    display: none !important;
+}
+.preset-grid label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+}
+#res-mode .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(3, 1fr) !important;
+    gap: 8px !important;
+}
+#res-mode label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
+#output-format .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(3, 1fr) !important;
+    gap: 8px !important;
+}
+#output-format label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
 #gpu-size .wrap {
     display: grid !important;
     grid-template-columns: repeat(2, 1fr) !important;
@@ -958,8 +1015,32 @@ css += """
     line-height: 1.5 !important;
     opacity: 0.7;
 }
+#res-auto-scale .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(4, minmax(0, 1fr)) !important;
+    gap: 6px !important;
+}
+#res-auto-scale label {
+    width: 100% !important;
+    min-width: 0 !important;
+    margin: 0 !important;
+    padding-left: 8px !important;
+    padding-right: 8px !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
 """
 
+_div_a = f'.preset-grid label:has(input[value="{PRESET_WIDE[0]}"])'
+_div_b = f'.preset-grid label:has(input[value="{PRESET_WIDE[1]}"])'
+css += (
+    f"{_div_a}, {_div_b} {{ position: relative !important; margin-top: 12px !important; overflow: visible !important; }}\n"
+    f"{_div_a}::before, {_div_b}::before {{ content: \"\"; position: absolute; top: -10px; height: 1px; "
+    f"background: rgba(128, 128, 128, 0.55); pointer-events: none; }}\n"
+    f"{_div_a}::before {{ left: 0; right: -4px; }}\n"
+    f"{_div_b}::before {{ left: -4px; right: 0; }}\n"
+)
+
 _theme_dir = os.path.dirname(os.path.abspath(__file__))
 _theme_pick = os.environ.get("UI_THEME", "")
 if _theme_pick == "off":
@@ -980,8 +1061,7 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
                 gr.HTML(
                     '<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;600&display=swap">'
                     '<div id="logo-title">'
-                    '<h1 id="ast-main-title">Multi Ratio Image Editor</h1>'
-                    '<br>'
+                    '<h1 id="ast-main-title">Qwen 2511 + FLUX.2</h1>'
                     '<h1 id="ast-nfaa-title">Not For All Audiences</h1>'
                     '<div id="ast-notice">'
                     '<p><b>Adults only:</b> Zero tolerance for any content involving minors. '
@@ -997,7 +1077,8 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
                     with gr.Column():
                         with gr.Row():
                             image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
-                            image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
+                            with gr.Column():
+                                image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
 
                         prompt = gr.Text(
                             label="Prompt 🪄",
@@ -1012,26 +1093,8 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
                         )
 
                         with gr.Group():
-                            aspect_ratio = gr.Radio(
-                                label="Aspect Ratio",
-                                show_label=True,
-                                choices=ASPECT_RATIOS,
-                                value=RATIO_DEFAULT,
-                                elem_id="aspect-ratio",
-                            )
-
-                        with gr.Group():
-                            gr.Markdown("Lanczos Upscale", elem_classes=["lanczos-title"])
-                            with gr.Row(elem_id="lanczos-row"):
-                                scale_basic = gr.Checkbox(label="Basic", value=False, min_width=0)
-                                scale_medium = gr.Checkbox(label="Medium", value=False, min_width=0)
-                                scale_high = gr.Checkbox(label="High", value=False, min_width=0)
-
-                        with gr.Group():
-                            gr.Markdown("Output Format", elem_classes=["lanczos-title"])
                             output_format = gr.Radio(
                                 label="Output Format",
-                                show_label=False,
                                 choices=["PNG", "JPG", "WebP"],
                                 value="PNG",
                                 elem_id="output-format",
@@ -1076,10 +1139,50 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
                                     elem_id="gpu-size-note",
                                 )
 
-                    with gr.Column():
-                        result = gr.Gallery(label="Result", show_label=False, type="filepath")
-                        use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
+                            with gr.Group():
+                                resolution_mode = gr.Radio(
+                                    label="Resolution mode",
+                                    choices=["Auto", "Preset", "Custom"],
+                                    value="Auto",
+                                    elem_id="res-mode",
+                                )
+                                resolution_preset = gr.Radio(
+                                    label="Resolution",
+                                    choices=PRESET_ALL,
+                                    value=PRESET_DEFAULT,
+                                    elem_classes=["preset-grid"],
+                                    elem_id="res-preset",
+                                )
+                                custom_height = gr.Slider(
+                                    label="Height", minimum=256, maximum=2048, step=8, value=256,
+                                    buttons=[], elem_id="res-custom-height",
+                                )
+                                custom_width = gr.Slider(
+                                    label="Width", minimum=256, maximum=2048, step=8, value=256,
+                                    buttons=[], elem_id="res-custom-width",
+                                )
+                                auto_scale = gr.Radio(
+                                    label="Auto resolution",
+                                    choices=AUTO_SCALES,
+                                    value="1x",
+                                    elem_id="res-auto-scale",
+                                )
 
+                    with gr.Column():
+                        result = gr.Gallery(label="Result", show_label=False, type="filepath", interactive=True, elem_id="result-gallery")
+                        with gr.Column(visible=False) as refine_group:
+                            with gr.Row():
+                                use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
+                                refine_with_flux_btn = gr.Button(
+                                    "Refine with FLUX.2-klein-9B", variant="primary", size="lg",
+                                )
+                            flux_seed = gr.Slider(label="FLUX seed", minimum=0, maximum=MAX_SEED, step=1, value=42, elem_id="flux-seed-slider")
+                            flux_randomize_seed = gr.Checkbox(label="Randomize FLUX seed", value=False)
+                        flux_result = gr.Gallery(label="FLUX Result", show_label=True, type="filepath", interactive=True, elem_id="flux-result-gallery")
+                        with gr.Row():
+                            use_flux_output_btn = gr.Button(
+                                "↗️ Use FLUX output as input", variant="secondary", size="sm", visible=False,
+                            )
 
         with gr.Tab("🔒", visible=False) as admin_tab:
             with gr.Group(visible=True) as admin_lock_group:
@@ -1136,23 +1239,9 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
                                         **_PRIVATE,
                                     )
 
-    for _box, _js in (
-        (scale_basic, "(on) => [on, false, false]"),
-        (scale_medium, "(on) => [false, on, false]"),
-        (scale_high, "(on) => [false, false, on]"),
-    ):
-        _box.input(
-            fn=None,
-            inputs=[_box],
-            outputs=[scale_basic, scale_medium, scale_high],
-            js=_js,
-            show_progress="hidden",
-            **_PRIVATE,
-        )
-
-    edit_event = gr.on(
+    gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=run_qwen_and_flux,
+        fn=infer_qwen_only,
         inputs=[
             image_1,
             image_2,
@@ -1163,28 +1252,54 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
             num_inference_steps,
             gpu_duration,
             gpu_size,
-            aspect_ratio,
-            scale_basic,
-            scale_medium,
-            scale_high,
+            resolution_mode,
+            resolution_preset,
+            custom_width,
+            custom_height,
+            auto_scale,
             output_format,
             age_confirm,
         ],
-        outputs=[result, seed, use_output_btn],
+        outputs=[result, seed, use_output_btn, refine_group],
         **({"show_progress_on": [result]} if _SHOW_PROGRESS_ON else {}),
-    )
-    edit_event.then(
+    ).then(
         fn=_sync_result_followups,
         inputs=result,
-        outputs=use_output_btn,
+        outputs=[use_output_btn, refine_group],
         show_progress="hidden",
     )
 
     use_output_btn.click(fn=use_output_as_input, inputs=[result], outputs=[image_1])
+
+    refine_with_flux_btn.click(
+        fn=refine_with_flux,
+        inputs=[result, flux_seed, flux_randomize_seed, gpu_duration, gpu_size, output_format],
+        outputs=[flux_result, flux_seed, use_flux_output_btn, refine_group],
+        **({"show_progress_on": [flux_result]} if _SHOW_PROGRESS_ON else {}),
+    ).then(
+        fn=_sync_result_followups,
+        inputs=result,
+        outputs=[use_output_btn, refine_group],
+        show_progress="hidden",
+    ).then(
+        fn=_sync_flux_result_followups,
+        inputs=flux_result,
+        outputs=use_flux_output_btn,
+        show_progress="hidden",
+    )
+
+    use_flux_output_btn.click(fn=use_output_as_input, inputs=[flux_result], outputs=[image_1])
+
+    flux_result.change(
+        fn=_sync_flux_result_followups,
+        inputs=flux_result,
+        outputs=use_flux_output_btn,
+        show_progress="hidden",
+    )
     result.change(
         fn=_sync_result_followups,
         inputs=result,
-        outputs=use_output_btn,
+        outputs=[use_output_btn, refine_group],
         show_progress="hidden",
     )
 
@@ -1206,6 +1321,45 @@ with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
         **_PRIVATE,
     )
 
+    result.upload(fn=lambda: None, inputs=None, outputs=result, show_progress="hidden", **_PRIVATE)
+    flux_result.upload(fn=lambda: None, inputs=None, outputs=flux_result, show_progress="hidden", **_PRIVATE)
+    demo.load(
+        fn=None, inputs=None, outputs=None,
+        js="""() => {
+            const ids = ["result-gallery", "flux-result-gallery"];
+            const inGallery = (t) => t && t.closest && ids.some((id) => t.closest("#" + id));
+            const lock = () => ids.forEach((id) => {
+                const root = document.getElementById(id);
+                if (!root) return;
+                root.querySelectorAll("input[type=file]").forEach((i) => {
+                    i.disabled = true;
+                    const wrap = i.closest("button");
+                    if (wrap) wrap.style.visibility = "hidden";
+                });
+            });
+            ["dragenter", "dragover", "drop"].forEach((ev) =>
+                document.addEventListener(ev, (e) => {
+                    if (inGallery(e.target)) { e.preventDefault(); e.stopPropagation(); }
+                }, true));
+            new MutationObserver(lock).observe(document.body, { childList: true, subtree: true });
+            lock();
+        }""",
+        **_PRIVATE,
+    )
+
+    resolution_mode.change(
+        fn=None, inputs=None, outputs=None,
+        js="""() => {
+            const nudge = () => {
+                window.dispatchEvent(new Event("resize"));
+                try { if (window.parentIFrame && window.parentIFrame.size) window.parentIFrame.size(); } catch (e) {}
+            };
+            [0, 100, 300, 800].forEach((ms) => setTimeout(nudge, ms));
+        }""",
+        show_progress="hidden",
+        **_PRIVATE,
+    )
+
     demo.load(fn=mark_browser_session, inputs=None, outputs=None, **_PRIVATE)
     demo.load(fn=check_admin_tab_visibility, inputs=None, outputs=[admin_tab], **_PRIVATE)