Add speed (1024px) / quality (2048px) toggle; decode 2K outputs with 1024px VAE tiles

app.py

@@ -34,6 +34,7 @@ PE_SPACE_ID = os.environ.get("PE_SPACE_ID", "hugging-apps/qwen-image-2-1-prompt-
 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.
@@ -186,6 +187,10 @@ def ratio_to_size(ratio, area=1024 * 1024, multiple=32):
     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)
+
+
+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
@@ -237,6 +242,13 @@ def generation_duration(prompt, image_paths, height, width, negative_prompt, see
 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(
@@ -285,7 +297,7 @@ logger = GenerationLogger(log_dir=LOG_DIR)
 MAX_SEED = np.iinfo(np.int32).max
 
 # --- Stage 1: screen and rewrite the prompt (CPU, no GPU quota) ---
-def prepare_stage(input_images, original_prompt, enable_extend, custom_size, seed, randomize_seed):
+def prepare_stage(input_images, original_prompt, enable_extend, custom_size, quality, seed, randomize_seed):
     """
     校验输入、NCII 检查（仅在有输入图片时）、可选的提示词改写，全部在 GPU 之外完成。
     Returns (image_paths, final_prompt, rewritten_prompt, seed, auto_height, auto_width).
@@ -310,8 +322,15 @@ def prepare_stage(input_images, original_prompt, enable_extend, custom_size, see
             gr.Warning("提示词改写暂不可用，已使用原始提示词。 / Prompt enhancement is unavailable; using the original prompt.")
 
     auto_height, auto_width = (None, None)
-    if not custom_size and not image_paths:
-        auto_height, auto_width = ratio_to_size(wh_ratio)
+    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:
+            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
 
 
@@ -399,9 +418,9 @@ def make_placeholder_image(text, width=512, height=320, bg_color=(30, 30, 30), t
 
 
 # --- Two-step flow: prepare (CPU) then generate (GPU) ---
-def prepare_request(input_images, original_prompt, enable_extend, custom_size, seed, randomize_seed):
+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, seed, randomize_seed,
+        input_images, original_prompt, enable_extend, custom_size, quality, seed, randomize_seed,
     )
     request_state = {
         "image_paths": image_paths,
@@ -576,6 +595,11 @@ with gr.Blocks(title="Qwen Image 2.1 Demo") as demo:
                 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(
@@ -720,6 +744,9 @@ with gr.Blocks(title="Qwen Image 2.1 Demo") 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"))
@@ -750,6 +777,7 @@ with gr.Blocks(title="Qwen Image 2.1 Demo") as demo:
             prompt,           # original_prompt
             enable_extend,    # enable_extend (是否开启提示词改写)
             custom_size,      # custom_size (是否自定义输出尺寸)
+            quality,
             seed,
             randomize_seed,
         ],