[Admin maintenance] Migrate to ZeroGPU (#2)

- [Admin maintenance] Migrate to ZeroGPU (970bb8eb02fb83fe1a4ea8742d765f77d68f7ed8)
- Run on a large ZeroGPU slice (f85981c3fc5e98fcb96ca7b75ab395a22fb137c0)
- Update app.py (91199a503ce1f7aafafff19b54a9a835c7af3b22)
- Add speed (1024px) / quality (2048px) toggle; decode 2K outputs with 1024px VAE tiles (e33ac7d1546107887d5e97b8d7972b79744e05c0)
- Show diffusion progress with gr.Progress(track_tqdm=True) (50995d8da0647cd2e0de08b4824ee16a06ba5ea1)


Co-authored-by: Apolinário from multimodal AI art <multimodalart@users.noreply.huggingface.co>

README.md

@@ -4,9 +4,10 @@ emoji: 🎨
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 5.33.0
+sdk_version: 6.28.0
 python_version: '3.10'
 app_file: app.py
+startup_duration_timeout: 1h
 pinned: true
 license: other
 license_name: qwen-research

app.py

@@ -1,31 +1,63 @@
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces
 import gradio as gr
 import numpy as np
 import random
-import requests
 import io
 import time
-# import spaces
 import uuid
 from datetime import datetime
 
 from PIL import Image
 
-import os
 import base64
 import json
+import re
+
+import torch
+from diffusers import QwenImage21Pipeline
+
+import ncii_guard
 
 # ============== 配置参数 ==============
 # 日志目录，可通过环境变量 LOG_DIR 自定义
 LOG_DIR = os.environ.get("LOG_DIR", "./generation_logs_paper_case")
 
-# ============== API 配置 ==============
-API_ENDPOINT = "https://poc-dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
-EDIT_MODEL = "pre-qwen-image-2.1-pro-yunqi"
-T2I_MODEL = "pre-qwen-image-2.1-pro-yunqi"
+# ============== 模型配置 ==============
+MODEL_ID = os.environ.get("QWEN_IMAGE_MODEL", "Qwen/Qwen-Image-2.1")
+# Prompt rewriting (Qwen-Image-2.1-PE-T2I / PE-I2I) runs in a companion Space so this
+# Space only holds the diffusion pipeline and fits a `large` ZeroGPU slice.
+PE_SPACE_ID = os.environ.get("PE_SPACE_ID", "hugging-apps/qwen-image-2-1-prompt-enhancer")
+PE_MAX_NEW_TOKENS = {"t2i": 1536, "i2i": 2048}
+GUARD_THRESHOLD = 0.5
+NUM_INFERENCE_STEPS = 28
+QUALITY_RESOLUTIONS = {"speed": 1024, "quality": 2048}
+TRUE_CFG_SCALE = 4.0
+# The prefix KV cache costs ~2 GB per 1K condition image; above this budget it is
+# switched off so many-image edits still fit next to the weights.
+KV_CACHE_BYTES_PER_TOKEN = 32 * 2 * 4096 * 2
+KV_CACHE_BUDGET_GB = 10.0
 
 MAX_INPUT_IMAGES = 10
 DEFAULT_LANGUAGE = "en"
 
+pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
+pipe.to("cuda")
+# Decode large outputs in tiles so a 2K VAE decode fits next to the weights.
+pipe.vae.enable_tiling(
+    tile_sample_min_height=1536,
+    tile_sample_min_width=1536,
+    tile_sample_stride_height=1152,
+    tile_sample_stride_width=1152,
+)
+
+# The NCII classifier runs in a CPU subprocess: a transformers forward in the main
+# process breaks the ZeroGPU worker fork.
+ncii_guard.start()
+
 # ============== 预设分辨率 ==============
 SIZE_PRESETS_2K = {
     "2688x1536 (16:9)": (2688, 1536),
@@ -140,295 +172,205 @@ def validate_image_count(images):
         raise gr.Error("最多支持 10 张输入图片。 / Up to 10 input images are supported.")
 
 
-def encode_image(pil_image):
-    buffered = io.BytesIO()
-    pil_image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+def check_prompt_guard(prompt):
+    """Reject image-editing prompts the NCII classifier flags. The error is
+    deliberately generic and does not say which classifier fired."""
+    if ncii_guard.score(prompt or "") >= GUARD_THRESHOLD:
+        raise gr.Error("prompt invalid based on our classifiers, try again")
 
 
-def call_image_api(prompt, mode="t2i", images=None, size="1024*1024", negative_prompt=" ", seed=None, prompt_extend=True):
-    """
-    调用 POC API 生成图片
-
-    Args:
-        prompt: 文本提示词
-        mode: "t2i" 或 "edit"
-        images: PIL Image 列表 (edit 模式需要)
-        size: 尺寸字符串如 "1024*1024"
-        negative_prompt: 负向提示词
-        seed: 随机种子
-        prompt_extend: 是否开启API侧提示词智能改写
-
-    Returns:
-        tuple: (PIL.Image, dict) 生成的图片和完整API响应
-    """
-    validate_image_count(images)
+def ratio_to_size(ratio, area=1024 * 1024, multiple=32):
+    """Turn a rewriter aspect ratio like "16:9" into a height/width pair of about `area` pixels."""
+    match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*:\s*(\d+(?:\.\d+)?)\s*", str(ratio or ""))
+    if not match or float(match.group(2)) == 0:
+        return None, None
+    aspect = float(match.group(1)) / float(match.group(2))
+    if not 1 / 4 <= aspect <= 4:
+        return None, None
+    return aspect_to_size(aspect, area, multiple)
 
-    # 使用独立的 IMAGE_API_KEY 环境变量
-    api_key = os.environ.get('IMAGE_API_KEY')
-    if not api_key:
-        raise EnvironmentError("IMAGE_API_KEY environment variable is not set")
 
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json"
-    }
+def aspect_to_size(aspect, area, multiple=32):
+    width = round((area * aspect) ** 0.5 / multiple) * multiple
+    height = round((area / aspect) ** 0.5 / multiple) * multiple
+    return height, width
 
-    # 构造 messages
-    content = [{"text": prompt}]
-
-    if mode == "edit":
-        model = EDIT_MODEL
-        # 添加图片到 content
-        if images and len(images) > 0:
-            # 图片放在文本之前
-            content = []
-            content.append({"text": prompt})
-            for img in images:
-                img_base64 = encode_image(img)
-                content.append({"image": f"data:image/png;base64,{img_base64}"})
-    else:
-        model = T2I_MODEL
-
-    # 构造请求体
-    payload = {
-        "model": model,
-        "input": {
-            "messages": [
-                {
-                    "role": "user",
-                    "content": content
-                }
-            ]
-        },
-        "parameters": {
-            "watermark": False,
-            "negative_prompt": negative_prompt,
-            "prompt_extend": prompt_extend,
-            "debug": True
-        }
-    }
-
-    # 传入 seed 参数
-    if seed is not None:
-        payload["parameters"]["seed"] = int(seed)
-
-    # 传入 size 参数（T2I 始终传，Edit 仅在指定时传）
-    if size:
-        payload["parameters"]["size"] = size
-
-    print(f"[API] Calling {model} API...")
-    print(f"[API] Prompt: {prompt[:100]}..." if len(prompt) > 100 else f"[API] Prompt: {prompt}")
-    if mode == "edit" and images:
-        print(f"[API] Input images count: {len(images)}")
 
-    # 发送请求
-    response = requests.post(API_ENDPOINT, headers=headers, json=payload)
-    print(response)
+_pe_client = None
 
-    if response.status_code != 200:
-        raise Exception(f"API request failed with status {response.status_code}: {response.text}")
 
-    result = response.json()
-    print(f"[API] Response: {json.dumps(result, ensure_ascii=False)[:500]}...")
+def enhance_prompt(prompt, image_paths):
+    """Rewrite the prompt with the official PE models in the companion Space.
+    Returns (rewritten_prompt, wh_ratio)."""
+    global _pe_client
+    from gradio_client import Client, handle_file
 
-    # 解析响应 - 处理同步和异步两种情况
-    if "output" in result:
-        output = result["output"]
-
-        # 检查是否是异步任务
-        if "task_id" in output and "task_status" in output:
-            task_id = output["task_id"]
-            task_status = output["task_status"]
-            print(f"[API] Async task created: {task_id}, status: {task_status}")
-
-            # 轮询等待任务完成
-            image, async_result = poll_task_result(task_id, api_key)
-            return image, async_result
+    if _pe_client is None:
+        _pe_client = Client(PE_SPACE_ID, token=os.environ.get("HF_TOKEN"), httpx_kwargs={"timeout": 900}, verbose=False)
+    rewritten, wh_ratio, *_ = _pe_client.predict(
+        prompt=prompt,
+        image_paths=[handle_file(p) for p in image_paths],
+        max_new_tokens=PE_MAX_NEW_TOKENS["i2i" if image_paths else "t2i"],
+        enable_thinking=False,
+        seed=0,
+        randomize_seed=True,
+        api_name="/enhance",
+    )
+    return str(rewritten or "").strip(), str(wh_ratio or "").strip()
+
+
+def kv_cache_fits(n_images):
+    return n_images * (1024 // 16) ** 2 * KV_CACHE_BYTES_PER_TOKEN <= KV_CACHE_BUDGET_GB * 1e9
+
+
+def generation_duration(prompt, image_paths, height, width, negative_prompt, seed):
+    """GPU seconds, fitted on this pipeline (eager, large ZeroGPU): step cost scales as
+    latent_tokens ** 1.363 and each cached condition image adds a fifth of its tokens."""
+    n_images = len(image_paths or [])
+    pixels = (int(width) * int(height)) if (width and height) else 1024 * 1024
+    prefix = n_images * (1024 // 16) ** 2
+    cached = kv_cache_fits(n_images)
+    tokens = pixels / 256 + (0.2 * prefix if cached else prefix)
+    per_step = 3.4e-6 * tokens ** 1.363 * 1.3
+    if negative_prompt:
+        per_step *= 2
+    fixed = 4 + 2 * n_images + 1.5e-6 * pixels
+    return int(min(300, (fixed + NUM_INFERENCE_STEPS * per_step) * 1.25))
+
+
+@spaces.GPU(duration=generation_duration)
+def run_pipeline(prompt, image_paths, height, width, negative_prompt, seed):
+    images = [Image.open(p) for p in image_paths] or None
+    kwargs = {"use_kv_cache": kv_cache_fits(len(image_paths))}
+    tile, stride = (1024, 768) if max(height or 0, width or 0) > 1536 else (1536, 1152)
+    pipe.vae.enable_tiling(
+        tile_sample_min_height=tile,
+        tile_sample_min_width=tile,
+        tile_sample_stride_height=stride,
+        tile_sample_stride_width=stride,
+    )
+    if negative_prompt:
+        kwargs.update(negative_prompt=negative_prompt, true_cfg_scale=TRUE_CFG_SCALE)
+    return pipe(
+        prompt,
+        image=images,
+        height=height,
+        width=width,
+        num_inference_steps=NUM_INFERENCE_STEPS,
+        generator=torch.Generator("cuda").manual_seed(int(seed)),
+        **kwargs,
+    ).images[0]
+
+
+def gallery_paths(input_images):
+    """Save gallery images to PNG files so they can be sent to the rewriter Space and
+    reopened in the GPU worker without loss (RGBA inputs keep their alpha)."""
+    import tempfile
+
+    paths = []
+    for item in input_images or []:
+        try:
+            if isinstance(item, (tuple, list)) and item:
+                item = item[0]
+            if isinstance(item, str):
+                item = Image.open(item)
+            elif hasattr(item, "name") and not isinstance(item, Image.Image):
+                item = Image.open(item.name)
+            if not isinstance(item, Image.Image):
+                continue
+            if item.mode not in ("RGB", "RGBA"):
+                item = item.convert("RGBA")
+            path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+            item.save(path)
+            paths.append(path)
+        except Exception as e:
+            print(f"[Warning] Failed to load input image: {e}")
+    return paths
 
-        # 同步响应 - 直接返回结果
-        if "choices" in output and len(output["choices"]) > 0:
-            choice = output["choices"][0]
-            if "message" in choice and "content" in choice["message"]:
-                for item in choice["message"]["content"]:
-                    if "image" in item:
-                        image_url = item["image"]
-                        return download_image(image_url), result
 
-        # 另一种响应格式
-        if "results" in output and len(output["results"]) > 0:
-            image_url = output["results"][0].get("url")
-            if image_url:
-                return download_image(image_url), result
+# 初始化日志记录器
+print(f"Initializing logger with directory: {LOG_DIR}")
+logger = GenerationLogger(log_dir=LOG_DIR)
 
-    raise Exception(f"Failed to parse API response: {result}")
 
+# --- UI Constants and Helpers ---
+MAX_SEED = np.iinfo(np.int32).max
 
-def poll_task_result(task_id, api_key, max_retries=120, interval=2):
+# --- Stage 1: screen and rewrite the prompt (CPU, no GPU quota) ---
+def prepare_stage(input_images, original_prompt, enable_extend, custom_size, quality, seed, randomize_seed):
     """
-    轮询异步任务结果
-
-    Args:
-        task_id: 任务ID
-        api_key: API密钥
-        max_retries: 最大重试次数
-        interval: 轮询间隔(秒)
-
-    Returns:
-        tuple: (PIL.Image, dict) 生成的图片和完整API响应
+    校验输入、NCII 检查（仅在有输入图片时）、可选的提示词改写，全部在 GPU 之外完成。
+    Returns (image_paths, final_prompt, rewritten_prompt, seed, auto_height, auto_width).
     """
-    task_url = f"https://poc-dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
-    headers = {
-        "Authorization": f"Bearer {api_key}"
-    }
-
-    for i in range(max_retries):
-        response = requests.get(task_url, headers=headers)
-        if response.status_code != 200:
-            print(f"[API] Task polling failed: {response.status_code}")
-            time.sleep(interval)
-            continue
-
-        result = response.json()
-        output = result.get("output", {})
-        task_status = output.get("task_status", "UNKNOWN")
-
-        print(f"[API] Task {task_id} status: {task_status} (attempt {i+1}/{max_retries})")
-
-        if task_status == "SUCCEEDED":
-            # 获取生成的图片
-            if "results" in output and len(output["results"]) > 0:
-                image_url = output["results"][0].get("url")
-                if image_url:
-                    return download_image(image_url), result
-            raise Exception(f"Task succeeded but no image found in response: {result}")
+    validate_image_count(input_images)
+    if not original_prompt or not original_prompt.strip():
+        raise gr.Error("请输入提示词。 / Please enter a prompt.")
+    image_paths = gallery_paths(input_images)
 
-        elif task_status == "FAILED":
-            error_msg = output.get("message", "Unknown error")
-            raise Exception(f"Task failed: {error_msg}")
+    if image_paths:
+        check_prompt_guard(original_prompt)
 
-        elif task_status in ["PENDING", "RUNNING"]:
-            time.sleep(interval)
-            continue
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
+    rewritten_prompt, wh_ratio = "", ""
+    if enable_extend:
+        try:
+            rewritten_prompt, wh_ratio = enhance_prompt(original_prompt, image_paths)
+        except Exception as e:
+            print(f"[Warning] Prompt enhancement failed: {e!r}")
+            gr.Warning("提示词改写暂不可用，已使用原始提示词。 / Prompt enhancement is unavailable; using the original prompt.")
+
+    auto_height, auto_width = (None, None)
+    if not custom_size:
+        area = QUALITY_RESOLUTIONS.get(quality, 1024) ** 2
+        if image_paths:
+            last_width, last_height = Image.open(image_paths[-1]).size
+            auto_height, auto_width = aspect_to_size(last_width / last_height, area)
         else:
-            print(f"[API] Unknown task status: {task_status}")
-            time.sleep(interval)
-
-    raise Exception(f"Task {task_id} timed out after {max_retries * interval} seconds")
-
-
-def download_image(url):
-    """
-    从URL下载图片
-
-    Args:
-        url: 图片URL
-
-    Returns:
-        PIL.Image: 下载的图片
-    """
-    print(f"[API] Downloading image from: {url[:100]}...")
-    response = requests.get(url)
-    if response.status_code != 200:
-        raise Exception(f"Failed to download image: {response.status_code}")
+            auto_height, auto_width = ratio_to_size(wh_ratio, area)
+            if auto_height is None:
+                auto_height, auto_width = aspect_to_size(1.0, area)
+    return image_paths, rewritten_prompt or original_prompt, rewritten_prompt, seed, auto_height, auto_width
 
-    with Image.open(io.BytesIO(response.content)) as image:
-        return image.convert("RGBA")
-
-
-
-
-# 初始化日志记录器
-print(f"Initializing logger with directory: {LOG_DIR}")
-logger = GenerationLogger(log_dir=LOG_DIR)
-
-
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
 
 # --- Stage 2: Image Generation Function ---
 def generate_image_stage(
-    input_images,
+    image_paths,
     original_prompt,
+    final_prompt,
     custom_size,
     log_dir,
-    seed=42,
-    randomize_seed=False,
-    height=1024,
-    width=1024,
+    seed,
+    height,
+    width,
+    auto_height,
+    auto_width,
     negative_prompt=" ",
     prompt_extend=True,
     username=None,
-    progress=gr.Progress(track_tqdm=True),
 ):
     """
-    使用 POC API 生成图片。
     根据是否有输入图片自动判断模式：有图片 -> Edit，无图片 -> T2I。
-    改写由 API 侧 prompt_extend 参数控制。
     """
-    validate_image_count(input_images)
-
     # 更新日志目录（如果用户修改了）
     if log_dir and log_dir.strip():
         logger.set_log_dir(log_dir.strip())
 
-    # 使用传入的 negative_prompt（默认为空格）
-    if not negative_prompt or not negative_prompt.strip():
-        negative_prompt = " "
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    # Load input images into PIL Images from gallery
-    pil_images = []
-    if input_images is not None and len(input_images) > 0:
-        for item in input_images:
-            try:
-                if isinstance(item, Image.Image):
-                    pil_images.append(item.convert("RGB"))
-                elif isinstance(item, tuple) and len(item) > 0 and isinstance(item[0], Image.Image):
-                    pil_images.append(item[0].convert("RGB"))
-                elif isinstance(item, str):
-                    pil_images.append(Image.open(item).convert("RGB"))
-                elif hasattr(item, "name"):
-                    pil_images.append(Image.open(item.name).convert("RGB"))
-            except Exception as e:
-                print(f"[Warning] Failed to load input image: {e}")
-                continue
-        print(f"[API] Loaded {len(pil_images)} input images for editing")
-
-    # 根据是否有输入图片自动判断模式
-    is_edit_mode = len(pil_images) > 0
+    negative_prompt = (negative_prompt or "").strip()
+    is_edit_mode = len(image_paths) > 0
     mode = "edit" if is_edit_mode else "t2i"
 
-    # 构造 size：仅在用户开启自定义尺寸时传，否则由 API 自动决定
+    # 构造 size：开启自定义尺寸时使用设置值，否则由改写模型推荐的比例或模型自动决定
     if custom_size:
-        size = f"{width}*{height}"
+        out_height, out_width = int(height), int(width)
     else:
-        size = None
-
-    print(f"[API] Mode: {mode} (auto-detected: {'has input images' if is_edit_mode else 'no input images'})")
-    print(f"[API] Prompt: '{original_prompt[:100]}...'" if len(original_prompt) > 100 else f"[API] Prompt: '{original_prompt}'")
-    print(f"[API] Negative Prompt: '{negative_prompt}'")
-    print(f"[API] Prompt Extend: {prompt_extend}")
-    print(f"[API] Input images count: {len(pil_images)}")
-    print(f"[API] Seed: {seed}, Size: {size if size else 'auto (not specified)'}")
-
-    # 调用 API 生成图片
-    image, api_response = call_image_api(
-        prompt=original_prompt,
-        mode=mode,
-        images=pil_images if len(pil_images) > 0 else None,
-        size=size,
-        negative_prompt=negative_prompt,
-        seed=seed,
-        prompt_extend=prompt_extend
-    )
+        out_height, out_width = auto_height, auto_width
+
+    print(f"Mode: {mode}, prompt extend: {prompt_extend}, input images: {len(image_paths)}, "
+          f"seed: {seed}, size: {f'{out_width}x{out_height}' if out_width else 'auto'}")
 
-    # Only show actual API rewrite text; never substitute the original prompt.
-    rewritten_prompt = extract_rewritten_prompt(api_response) if prompt_extend else ""
-    enhanced_prompt = rewritten_prompt or original_prompt
+    image = run_pipeline(final_prompt, image_paths, out_height, out_width, negative_prompt, seed)
 
     # 记录生成日志
     params = {
@@ -437,50 +379,25 @@ def generate_image_stage(
         "width": width,
         "negative_prompt": negative_prompt,
         "prompt_extend": prompt_extend,
-        "input_images_count": len(pil_images),
-        "api_model": EDIT_MODEL if is_edit_mode else T2I_MODEL
+        "input_images_count": len(image_paths),
+        "model": MODEL_ID,
+        "num_inference_steps": NUM_INFERENCE_STEPS,
     }
 
     log_name = logger.log_generation(
         original_prompt=original_prompt,
-        enhanced_prompt=enhanced_prompt,
+        enhanced_prompt=final_prompt,
         image=image,
         seed=seed,
         params=params,
         gpu_id=None,
-        input_images=pil_images if len(pil_images) > 0 else None,
+        input_images=[Image.open(p) for p in image_paths] or None,
         username=username,
-        api_response=api_response
     )
 
-    print(f"[API] Generation complete, logged as: {log_name}")
-
-    return image, seed, rewritten_prompt
-
-
-def extract_rewritten_prompt(api_response):
-    """Read the text returned alongside the image when parameters.debug is true."""
-    if not isinstance(api_response, dict):
-        return ""
-    output = api_response.get("output") or {}
-    debug_info = output.get("debug_info") or {}
-    rewrite_info = debug_info.get("rewrite_debug_info") or {}
-    rewritten = rewrite_info.get("rewritten_prompt")
-    if isinstance(rewritten, str) and rewritten.strip():
-        return rewritten
-    # Some responses expose only the prompt actually passed to generation.
-    actual_prompt = debug_info.get("actual_prompt")
-    if output.get("rewrite_status") == "success" and isinstance(actual_prompt, str) and actual_prompt.strip():
-        return actual_prompt
-    # Compatibility with responses that include rewrite text beside the image.
-    choices = output.get("choices") or []
-    if not choices:
-        return ""
-    content = (choices[0].get("message") or {}).get("content") or []
-    return "\n\n".join(
-        item["text"] for item in content
-        if isinstance(item, dict) and isinstance(item.get("text"), str) and item["text"].strip()
-    )
+    print(f"Generation complete, logged as: {log_name}")
+
+    return image
 
 
 def make_placeholder_image(text, width=512, height=320, bg_color=(30, 30, 30), text_color=(200, 200, 200)):
@@ -500,37 +417,48 @@ def make_placeholder_image(text, width=512, height=320, bg_color=(30, 30, 30), t
     return img
 
 
-# --- Unified generate_with_enhance generator ---
-def generate_with_enhance(
-    input_images,
+# --- Two-step flow: prepare (CPU) then generate (GPU) ---
+def prepare_request(input_images, original_prompt, enable_extend, custom_size, quality, seed, randomize_seed):
+    image_paths, final_prompt, rewritten_prompt, seed, auto_height, auto_width = prepare_stage(
+        input_images, original_prompt, enable_extend, custom_size, quality, seed, randomize_seed,
+    )
+    request_state = {
+        "image_paths": image_paths,
+        "final_prompt": final_prompt,
+        "auto_height": auto_height,
+        "auto_width": auto_width,
+    }
+    return make_placeholder_image("Generating image..."), seed, rewritten_prompt, request_state
+
+
+def generate_request(
+    request_state,
     original_prompt,
     enable_extend,
     custom_size,
     log_dir,
     seed,
-    randomize_seed,
     height,
     width,
     negative_prompt,
     request: gr.Request,
+    progress=gr.Progress(track_tqdm=True),
 ):
     """
-    生成图片，prompt_extend 由 API 侧处理。
-    Yields intermediate results so the UI updates progressively.
+    生成图片。提示词改写已在上一步完成（开启时）。
     """
+    if not request_state:
+        raise gr.Error("请重新点击生成。 / Please click generate again.")
     username = request.username if request else None
-    print(f"[API] Request from user: {username}")
-
-    yield make_placeholder_image("Generating image..."), seed, ""
-
-    image, seed, rewritten_prompt = generate_image_stage(
-        input_images, original_prompt,
-        custom_size, log_dir, seed, randomize_seed, height, width,
+    print(f"Request from user: {username}")
+    return generate_image_stage(
+        request_state["image_paths"], original_prompt, request_state["final_prompt"],
+        custom_size, log_dir, seed, height, width,
+        request_state["auto_height"], request_state["auto_width"],
         negative_prompt=negative_prompt,
         prompt_extend=enable_extend,
         username=username,
     )
-    yield image, seed, rewritten_prompt
 
 
 def make_example_loader(images, text, extend):
@@ -627,14 +555,14 @@ css = """
 #edit_text{margin-top: -62px !important}
 """
 
-with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
+with gr.Blocks(title="Qwen Image 2.1 Demo") as demo:
     with gr.Column(elem_id="col-container"):
         gr.HTML('<a href="https://huggingface.co/Qwen/Qwen-Image-2.1" target="_blank"><img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/image2.1/logo.png" alt="Qwen-Image Logo" width="400" style="display: block; margin: 0 auto;"></a>')
         language = gr.Radio(choices=[("中文", "zh"), ("English", "en")], value=DEFAULT_LANGUAGE, label="语言 / Language")
         instructions = gr.Markdown("""
 ## Qwen Image 2.1 Demo 使用说明
 1. 如果不输入图片，默认进入文生图模式；如果输入图片，则进入图像编辑模式（支持1-10张图片）。
-2. 默认"Enable Prompt Extend"为开启状态，API会自动进行提示词智能改写。如果希望直接使用原始提示词，可以关闭该选项。
+2. 默认"Enable Prompt Extend"为开启状态，会自动进行提示词智能改写。如果希望直接使用原始提示词，可以关闭该选项。
 3. 下方提供了包括文生图，图���图的测试样例，可以作为模型基础能力的参考。
 4. 生成透明图时，建议提示词遵循以下格式，将中间的省略号替换为具体画面描述：
 
@@ -665,16 +593,21 @@ with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
             )
         with gr.Row():
             enable_extend = gr.Checkbox(
-                label="Enable Prompt Extend (API侧提示词智能改写)",
+                label="Enable Prompt Extend (提示词智能改写)",
                 value=True,
             )
+            quality = gr.Radio(
+                choices=[("Speed (1024px)", "speed"), ("Quality (2048px)", "quality")],
+                value="speed",
+                show_label=False,
+            )
             generate_button = gr.Button("Generate Image", variant="primary")
 
         rewritten_prompt_output = localize(
             gr.Textbox(value="", lines=4, max_lines=20, interactive=False),
             label=("改写结果", "Rewritten prompt"),
-            placeholder=("开启智能改写并生成图片后，API 返回的改写提示词会显示在这里；未返回时留空。",
-                         "After generation with prompt enhancement enabled, the API rewrite appears here. Blank if none is returned."),
+            placeholder=("开启智能改写并生成图片后，改写后的提示词会显示在这里。",
+                         "After generation with prompt enhancement enabled, the rewritten prompt appears here."),
         )
         enable_extend.change(fn=lambda: "", inputs=[], outputs=[rewritten_prompt_output], queue=False)
 
@@ -691,8 +624,6 @@ with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
                 visible=False,
             )
 
-            # gr.Markdown(f"**当前使用 API 模式** (Edit: {EDIT_MODEL}, T2I: {T2I_MODEL})")
-
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -715,7 +646,7 @@ with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
             custom_size = gr.Checkbox(
                 label="自定义输出尺寸 (Customize output size)",
                 value=False,
-                info="关闭时由 API/模型自动决定输出尺寸；开启后使用下方分辨率设置",
+                info="关闭时由模型自动决定输出尺寸；开启后使用下方分辨率设置",
             )
 
             # 分辨率预设选择
@@ -814,6 +745,9 @@ with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
     localize(result, label=("生成结果", "Result"))
     localize(prompt, label=("提示词", "Prompt"), placeholder=("描述想生成或编辑的内容，可按上传顺序引用第 1–10 张图…", "Describe what to generate or edit; refer to images 1–10 in upload order…"))
     localize(enable_extend, label=("智能改写提示词", "Enhance prompt"))
+    localize(quality, choices=(
+        [("速度 (1024px)", "speed"), ("质量 (2048px)", "quality")],
+        [("Speed (1024px)", "speed"), ("Quality (2048px)", "quality")]))
     localize(generate_button, value=("生成图片", "Generate image"))
     localize(advanced, label=("高级设置", "Advanced settings"))
     localize(log_dir_input, label=("日志保存目录", "Log directory"), placeholder=("输入日志保存目录", "Enter a log directory"))
@@ -834,27 +768,39 @@ with gr.Blocks(title="Qwen Image 2.1 Demo", css=css) as demo:
     language.change(fn=switch_language, inputs=[language], outputs=localized_components, queue=False)
     input_images.upload(fn=validate_image_count, inputs=[input_images], outputs=[], queue=False)
 
-    # Generate Image button event (with optional prompt enhancement)
+    # Generate Image button event: screen/rewrite the prompt off-GPU, then generate.
+    # Two events so the GPU step is scheduled with a fresh ZeroGPU token after a long rewrite.
+    request_state = gr.State(None)
     generate_button.click(
-        fn=generate_with_enhance,
+        fn=prepare_request,
         inputs=[
             input_images,     # input_images (有图片->Edit模式, 无图片->T2I模式)
             prompt,           # original_prompt
-            enable_extend,    # enable_extend (是否开启API侧改写)
+            enable_extend,    # enable_extend (是否开启提示词改写)
             custom_size,      # custom_size (是否自定义输出尺寸)
-            log_dir_input,    # log_dir
+            quality,
             seed,
             randomize_seed,
+        ],
+        outputs=[result, seed, rewritten_prompt_output, request_state],
+        concurrency_limit=2,
+    ).success(
+        fn=generate_request,
+        inputs=[
+            request_state,
+            prompt,
+            enable_extend,
+            custom_size,
+            log_dir_input,    # log_dir
+            seed,
             height,
             width,
             negative_prompt_input,  # negative_prompt (负向提示词)
         ],
-        outputs=[result, seed, rewritten_prompt_output],
-        concurrency_limit=2
+        outputs=[result],
+        concurrency_limit=2,
     )
 
-demo.queue(default_concurrency_limit=10, max_size=20)
-
 if __name__ == "__main__":
     # 使用 os.path.realpath 解析真实路径，避免 NAS 软链接导致 Gradio 文件权限检查失败（404）
     _script_dir = os.path.realpath(os.path.dirname(os.path.abspath(__file__)))
@@ -865,4 +811,5 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         allowed_paths=_allowed,
+        css=css,
     )

ncii_guard.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import json
+import os
+import select
+import subprocess
+import sys
+import tempfile
+import threading
+
+GUARD_ID = "hfmlsoc/ncii-guard-v02"
+
+_lock = threading.Lock()
+_process: subprocess.Popen | None = None
+
+
+def _read(timeout: float) -> dict:
+    readable, _, _ = select.select([_process.stdout], [], [], timeout)
+    if not readable:
+        raise TimeoutError(f"the guard did not answer within {timeout}s")
+    line = _process.stdout.readline()
+    if not line:
+        raise EOFError("the guard process died")
+    return json.loads(line)
+
+
+def _spawn() -> None:
+    global _process
+    _process = subprocess.Popen(
+        [sys.executable, os.path.abspath(__file__)],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+    )
+    ready = _read(600.0)
+    if ready.get("status") != "ready":
+        raise RuntimeError(f"guard failed to start: {ready}")
+
+
+def start() -> None:
+    with _lock:
+        _spawn()
+
+
+def score(prompt: str, timeout: float = 60.0) -> float:
+    with _lock:
+        for attempt in (0, 1):
+            try:
+                if _process is None or _process.poll() is not None:
+                    _spawn()
+                _process.stdin.write(json.dumps({"prompt": prompt or ""}) + "\n")
+                _process.stdin.flush()
+                return float(_read(timeout)["score"])
+            except Exception:
+                if attempt:
+                    raise
+                if _process is not None and _process.poll() is None:
+                    _process.kill()
+
+
+def _snapshot() -> str:
+    from huggingface_hub import snapshot_download
+
+    src = snapshot_download(
+        GUARD_ID,
+        allow_patterns=["config.json", "tokenizer.json", "tokenizer_config.json", "model.safetensors"],
+    )
+    dst = os.path.join(tempfile.gettempdir(), "ncii-guard-v02-normalised")
+    os.makedirs(dst, exist_ok=True)
+    for name in os.listdir(src):
+        link = os.path.join(dst, name)
+        if not os.path.exists(link):
+            os.symlink(os.path.realpath(os.path.join(src, name)), link)
+
+    config = json.load(open(os.path.join(src, "config.json")))
+    rope = config.get("rope_parameters")
+    if isinstance(rope, dict):
+        per_layer = {k: v for k, v in rope.items() if isinstance(v, dict)}
+        if per_layer:
+            config["rope_parameters"] = per_layer
+    config_path = os.path.join(dst, "config.json")
+    if os.path.islink(config_path):
+        os.remove(config_path)
+    with open(config_path, "w") as fh:
+        json.dump(config, fh)
+    return dst
+
+
+def _serve() -> None:
+    protocol = os.fdopen(os.dup(1), "w", buffering=1)
+    os.dup2(2, 1)
+
+    import torch
+    from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+    guard_dir = _snapshot()
+    tokenizer = AutoTokenizer.from_pretrained(guard_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(guard_dir, dtype=torch.float32).eval()
+    protocol.write(json.dumps({"status": "ready", "id2label": model.config.id2label}) + "\n")
+    for line in sys.stdin:
+        text = json.loads(line)["prompt"]
+        batch = tokenizer(text, truncation=True, max_length=256, padding=True, return_tensors="pt")
+        with torch.no_grad():
+            logits = model(**batch).logits.float()
+        protocol.write(json.dumps({"score": torch.softmax(logits, dim=-1)[0, 1].item()}) + "\n")
+
+
+if __name__ == "__main__":
+    _serve()

requirements.txt

@@ -0,0 +1,6 @@
+torchvision
+diffusers @ git+https://github.com/huggingface/diffusers.git
+transformers @ git+https://github.com/huggingface/transformers.git
+accelerate
+safetensors
+pillow