Update app.py

app.py

@@ -11,13 +11,18 @@ from PIL import Image
 
 from diffusers import QwenImage21Pipeline
 
+
 MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen-Image-2.1")
 MAX_SEED = 2**31 - 1
 MAX_REFERENCE_IMAGES = 10
+MAX_DURATION = 195
+
+# Qwen Image's latent-space geometry uses 16x spatial compression.
+SIZE_MULTIPLE = 16
+MAX_EDIT_LONG_SIDE = 2048
+
 
-# Aspect presets. The 2K sizes are the ones published on the Qwen/Qwen-Image-2.1
-# model card; the 1K row is the same aspect at a quarter of the pixels, which is
-# what keeps a ZeroGPU request inside a sane duration.
+# Text-to-image output presets.
 SIZE_PRESETS = {
     "1:1 — 1024x1024": (1024, 1024),
     "16:9 — 1376x768": (1376, 768),
@@ -36,16 +41,28 @@ SIZE_PRESETS = {
 }
 DEFAULT_PRESET = "1:1 — 2048x2048 (2K)"
 
+
+# Image-to-image output settings.
+EDIT_SIZE_PRESETS = {
+    "Square — 2048x2048": "square_2k",
+    "Preserve first image aspect — long side 2048": "long_side_2k",
+}
+DEFAULT_EDIT_SIZE = "Preserve first image aspect — long side 2048"
+
+
 # The model card's recommended wrapper for native transparent output.
 RGBA_PREFIX = "This is an RGBA image with transparency. "
 RGBA_SUFFIX = " The image has alpha channel and the background is transparent."
 
-pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
 
-# The decoder upcasts its resample layers to fp32, so a full-resolution 2K decode
-# allocates enough to take the ZeroGPU worker down. Tiling bounds that to one tile
-# regardless of output size; 1024 is large enough that 1K outputs still decode in a
-# single pass and never pay for blending.
+pipe = QwenImage21Pipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+
+# Avoid a large full-frame VAE decode. At 1024px, ordinary 1K images still decode
+# in one tile, while larger outputs use tiled decoding and require less peak VRAM.
 pipe.vae.enable_tiling(
     tile_sample_min_height=1024,
     tile_sample_min_width=1024,
@@ -61,67 +78,182 @@ def _resolve_seed(seed: int, randomize: bool) -> int:
 
 
 def _save_png(image: Image.Image) -> str:
-    # Named temp file, never a fixed path: handlers run concurrently.
+    """Save an image as a unique PNG file and return its path."""
     fd, path = tempfile.mkstemp(suffix=".png")
     os.close(fd)
     image.save(path, format="PNG")
     return path
 
 
-# ZeroGPU scales the requested duration by ~1.5x before checking it against the
-# visitor's tier cap, so anything above ~195s here is rejected outright on a free
-# account. Keep the estimate tight: a large request also loses queue priority and
-# trips the quota check for visitors with little time left.
-MAX_DURATION = 195
+def _round_to_multiple(value: float, multiple: int = SIZE_MULTIPLE) -> int:
+    """Round a dimension to the nearest valid Qwen latent-grid multiple."""
+    return max(multiple, int(round(value / multiple)) * multiple)
+
+
+def _edit_output_size(
+    first_image: Image.Image,
+    edit_size: str = DEFAULT_EDIT_SIZE,
+) -> tuple[int, int]:
+    """Determine the target canvas for an image-edit request.
+
+    Square mode always produces 2048x2048.
+
+    Long-side mode preserves the first reference image's aspect ratio, scales its
+    longest side to 2048 px, and rounds both dimensions to multiples of 16.
+    """
+    mode = EDIT_SIZE_PRESETS.get(
+        edit_size,
+        EDIT_SIZE_PRESETS[DEFAULT_EDIT_SIZE],
+    )
+
+    if mode == "square_2k":
+        return 2048, 2048
+
+    source_width, source_height = first_image.size
+    if source_width <= 0 or source_height <= 0:
+        raise gr.Error("The first reference image has invalid dimensions.")
+
+    scale = MAX_EDIT_LONG_SIDE / max(source_width, source_height)
+
+    width = _round_to_multiple(source_width * scale)
+    height = _round_to_multiple(source_height * scale)
+
+    # Rounding upward can make an edge 16 pixels too large. Scale down once more
+    # so that "long side 2048" always remains a real upper bound.
+    if max(width, height) > MAX_EDIT_LONG_SIDE:
+        second_scale = MAX_EDIT_LONG_SIDE / max(width, height)
+        width = _round_to_multiple(width * second_scale)
+        height = _round_to_multiple(height * second_scale)
+
+    # This is a final guard against edge cases from rounding.
+    width = min(width, MAX_EDIT_LONG_SIDE)
+    height = min(height, MAX_EDIT_LONG_SIDE)
+
+    return width, height
+
 
-# Latent tokens the transformer attends over, after the pipeline's 16x spatial
-# compression. Condition images are resized to the output resolution, so each one
-# adds about as many tokens as the output itself.
 def _tokens(width: int, height: int) -> int:
+    """Number of spatial latent tokens after 16x compression."""
     return (width // 16) * (height // 16)
 
 
-def _raw_budget(width: int, height: int, steps: int, references: int = 0) -> int:
-    """Seconds one request needs on ZeroGPU, before clamping.
+def _raw_budget(
+    width: int,
+    height: int,
+    steps: int,
+    references: int = 0,
+) -> int:
+    """Estimate the un-clamped ZeroGPU duration requirement.
 
-    Only the output tokens are queries; reference tokens join as keys and values
-    and are KV-cached after the first step, so attention costs O x (O + R) rather
-    than (O + R) squared. That asymmetry is why ten reference images are far
-    cheaper than a 2K output. Constants fitted on timings measured here:
-    1024x1024/30 steps -> 17s, 1376x768/30 -> 19s, one reference/24 -> 19s,
-    2048x2048/30 -> ~95s.
+    The output latents form attention queries. Reference-image latents are
+    key/value context and are cached after the first step. This makes the
+    approximate attention cost O x (O + R), not (O + R)^2.
     """
     out_tokens = _tokens(width, height)
     ref_tokens = out_tokens * references
-    per_step = 9.11e-5 * out_tokens + 5.78e-9 * out_tokens * (out_tokens + ref_tokens)
+
+    per_step = (
+        9.11e-5 * out_tokens
+        + 5.78e-9 * out_tokens * (out_tokens + ref_tokens)
+    )
     seconds = 12 + 4 * references + steps * per_step
+
     return int(max(40, seconds * 1.3))
 
 
-def _budget(width: int, height: int, steps: int, references: int = 0) -> int:
-    return min(MAX_DURATION, _raw_budget(width, height, steps, references))
+def _budget(
+    width: int,
+    height: int,
+    steps: int,
+    references: int = 0,
+) -> int:
+    """Return a reservation duration no larger than ZeroGPU permits."""
+    return min(
+        MAX_DURATION,
+        _raw_budget(width, height, steps, references),
+    )
+
 
+def _require_budget(
+    width: int,
+    height: int,
+    steps: int,
+    references: int = 0,
+) -> None:
+    """Reject work that cannot fit inside a single allowed ZeroGPU reservation."""
+    required = _raw_budget(width, height, steps, references)
 
-def _require_budget(width: int, height: int, steps: int, references: int = 0) -> None:
-    """Refuse a request that cannot finish inside the reservation we are allowed."""
-    if _raw_budget(width, height, steps, references) > MAX_DURATION:
+    if required > MAX_DURATION:
         raise gr.Error(
-            "This request needs more GPU time than a single ZeroGPU reservation allows. "
-            "Use fewer steps, fewer reference images, or a smaller size."
+            "This request needs more GPU time than a single ZeroGPU reservation "
+            "allows. Use fewer steps, fewer reference images, or a smaller size."
         )
 
 
-# Gradio hands these the same arguments as the handler, and gr.Examples supplies
-# only the columns it lists — so every parameter needs a default, and the extras
-# Gradio appends positionally (progress) are swallowed.
-def _t2i_duration(prompt=None, transparent=False, size_preset=DEFAULT_PRESET, steps=30, *args, **kwargs):
-    width, height = SIZE_PRESETS.get(size_preset, SIZE_PRESETS[DEFAULT_PRESET])
+def _gallery_item_path(item) -> str:
+    """Extract a filepath from Gradio Gallery's supported return formats."""
+    if isinstance(item, (tuple, list)):
+        return item[0]
+    return item
+
+
+def _t2i_duration(
+    prompt=None,
+    transparent=False,
+    size_preset=DEFAULT_PRESET,
+    steps=30,
+    *args,
+    **kwargs,
+):
+    width, height = SIZE_PRESETS.get(
+        size_preset,
+        SIZE_PRESETS[DEFAULT_PRESET],
+    )
     return _budget(width, height, int(steps or 30))
 
 
-def _edit_duration(instruction=None, reference_images=None, steps=30, *args, **kwargs):
-    count = len(reference_images or [])
-    return _budget(1024, 1024, int(steps or 30), references=count)
+def _edit_duration(
+    instruction=None,
+    reference_images=None,
+    edit_size=DEFAULT_EDIT_SIZE,
+    steps=30,
+    *args,
+    **kwargs,
+):
+    """Estimate image-edit duration from the actual selected output geometry."""
+    references = reference_images or []
+    count = len(references)
+
+    # Before an image is uploaded, reserve enough for the selected 2K operation.
+    if not references:
+        return _budget(
+            2048,
+            2048,
+            int(steps or 30),
+            references=0,
+        )
+
+    try:
+        first_path = _gallery_item_path(references[0])
+
+        with Image.open(first_path) as first_image:
+            width, height = _edit_output_size(first_image, edit_size)
+
+        return _budget(
+            width,
+            height,
+            int(steps or 30),
+            references=count,
+        )
+    except Exception:
+        # The browser may be mid-upload when the duration callback is evaluated.
+        # Use a safe upper-bound fallback in that situation.
+        return _budget(
+            2048,
+            2048,
+            int(steps or 30),
+            references=count,
+        )
 
 
 @spaces.GPU(duration=_t2i_duration)
@@ -136,21 +268,7 @@ def generate(
     randomize_seed: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate an image from a text prompt with Qwen-Image-2.1.
-
-    Args:
-        prompt: What to draw. English and Chinese both work; the model is strong at rendering text inside the image.
-        transparent: Ask for a native RGBA result with a transparent background, using the prompt format the model was trained on.
-        size_preset: Output aspect ratio and resolution. The "(2K)" rows are the sizes published on the model card and are noticeably slower.
-        steps: Number of denoising steps. 30 is a good default; 40 is the model card's setting.
-        true_cfg_scale: Classifier-free guidance. Qwen-Image-2.1 is meant to be sampled at 1.0 (no guidance); raise it only to make the negative prompt take effect.
-        negative_prompt: What to avoid. Ignored unless true_cfg_scale is above 1.0.
-        seed: Seed for reproducible results, used when randomize_seed is off.
-        randomize_seed: Draw a fresh random seed for this run.
-
-    Returns:
-        The generated image, a PNG file of it (alpha preserved), the seed that was used, and how long generation took.
-    """
+    """Generate an image from a text prompt."""
     if not prompt or not prompt.strip():
         raise gr.Error("Enter a prompt first.")
 
@@ -158,7 +276,11 @@ def generate(
     if transparent:
         text = RGBA_PREFIX + text + RGBA_SUFFIX
 
-    width, height = SIZE_PRESETS.get(size_preset, SIZE_PRESETS[DEFAULT_PRESET])
+    width, height = SIZE_PRESETS.get(
+        size_preset,
+        SIZE_PRESETS[DEFAULT_PRESET],
+    )
+
     _require_budget(width, height, int(steps))
     used_seed = _resolve_seed(seed, randomize_seed)
 
@@ -170,23 +292,36 @@ def generate(
         true_cfg_scale=float(true_cfg_scale),
         generator=torch.Generator("cuda").manual_seed(used_seed),
     )
-    if float(true_cfg_scale) > 1.0 and negative_prompt and negative_prompt.strip():
+
+    if (
+        float(true_cfg_scale) > 1.0
+        and negative_prompt
+        and negative_prompt.strip()
+    ):
         kwargs["negative_prompt"] = negative_prompt.strip()
 
     started = time.perf_counter()
     image = pipe(**kwargs).images[0]
     elapsed = time.perf_counter() - started
 
+    # An opaque text-to-image result is exported as RGB PNG. Native transparent
+    # results retain their RGBA alpha channel.
     if not transparent and image.mode == "RGBA":
         image = image.convert("RGB")
 
-    return image, _save_png(image), used_seed, f"{elapsed:.1f}s on {width}x{height}, {int(steps)} steps"
+    return (
+        image,
+        _save_png(image),
+        used_seed,
+        f"{elapsed:.1f}s on {width}x{height}, {int(steps)} steps",
+    )
 
 
 @spaces.GPU(duration=_edit_duration)
 def edit(
     instruction: str,
     reference_images: list,
+    edit_size: str = DEFAULT_EDIT_SIZE,
     steps: int = 30,
     true_cfg_scale: float = 1.0,
     negative_prompt: str = "",
@@ -194,55 +329,74 @@ def edit(
     randomize_seed: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Edit or recombine up to 10 reference images with a natural-language instruction.
-
-    Args:
-        instruction: What to change, e.g. "change the background to a sunset beach" or "put the person from image 1 in the room from image 2".
-        reference_images: One to ten condition images. The model reads them as vision context, so identity and layout carry across.
-        steps: Number of denoising steps. 30 is a good default; 40 is the model card's setting.
-        true_cfg_scale: Classifier-free guidance. Leave at 1.0 unless you need the negative prompt to apply.
-        negative_prompt: What to avoid. Ignored unless true_cfg_scale is above 1.0.
-        seed: Seed for reproducible results, used when randomize_seed is off.
-        randomize_seed: Draw a fresh random seed for this run.
-
-    Returns:
-        The edited image, a PNG file of it (alpha preserved), the seed that was used, and how long the edit took.
-    """
+    """Edit or recombine up to ten images using a natural-language instruction."""
     if not instruction or not instruction.strip():
         raise gr.Error("Describe the edit you want.")
+
     if not reference_images:
         raise gr.Error("Add at least one reference image.")
+
     if len(reference_images) > MAX_REFERENCE_IMAGES:
         raise gr.Error(
             f"Up to {MAX_REFERENCE_IMAGES} reference images are supported."
         )
 
     images = []
+
     for item in reference_images:
-        path = item[0] if isinstance(item, (tuple, list)) else item
-        images.append(Image.open(path))
+        path = _gallery_item_path(item)
+
+        # Copy while the input file is open. This closes file handles before GPU
+        # inference and avoids lazy reads from temporary Gradio upload paths.
+        with Image.open(path) as source:
+            images.append(source.copy())
+
+    # The first reference defines the aspect ratio for long-side mode.
+    width, height = _edit_output_size(images[0], edit_size)
+
+    _require_budget(
+        width,
+        height,
+        int(steps),
+        references=len(images),
+    )
 
-    _require_budget(1024, 1024, int(steps), references=len(images))
     used_seed = _resolve_seed(seed, randomize_seed)
+
     kwargs = dict(
         prompt=instruction.strip(),
         image=images,
+        width=width,
+        height=height,
         num_inference_steps=int(steps),
         true_cfg_scale=float(true_cfg_scale),
         generator=torch.Generator("cuda").manual_seed(used_seed),
     )
-    if float(true_cfg_scale) > 1.0 and negative_prompt and negative_prompt.strip():
+
+    if (
+        float(true_cfg_scale) > 1.0
+        and negative_prompt
+        and negative_prompt.strip()
+    ):
         kwargs["negative_prompt"] = negative_prompt.strip()
 
     started = time.perf_counter()
     image = pipe(**kwargs).images[0]
     elapsed = time.perf_counter() - started
 
-    return image, _save_png(image), used_seed, f"{elapsed:.1f}s from {len(images)} reference image(s), {int(steps)} steps"
+    # _save_png preserves RGBA alpha when the pipeline supplies it.
+    return (
+        image,
+        _save_png(image),
+        used_seed,
+        f"{elapsed:.1f}s from {len(images)} reference image(s), "
+        f"{width}x{height}, {int(steps)} steps",
+    )
 
 
 INTRO = """ """
 
+
 with gr.Blocks(title=" ") as demo:
     gr.Markdown(INTRO)
 
@@ -254,44 +408,126 @@ with gr.Blocks(title=" ") as demo:
                         label=" ",
                         lines=10,
                     )
+
                     t2i_transparent = gr.Checkbox(
                         label=" ",
                         value=False,
                     )
+
                     t2i_size = gr.Dropdown(
                         label=" ",
                         choices=list(SIZE_PRESETS),
                         value=DEFAULT_PRESET,
                     )
-                    t2i_run = gr.Button(" ", variant="primary")
+
+                    t2i_run = gr.Button(
+                        " ",
+                        variant="primary",
+                    )
 
                     with gr.Accordion(" ", open=False):
-                        t2i_steps = gr.Slider(10, 50, value=25, step=1, label=" ")
+                        t2i_steps = gr.Slider(
+                            10,
+                            50,
+                            value=25,
+                            step=1,
+                            label=" ",
+                        )
+
                         t2i_cfg = gr.Slider(
-                            1.0, 10.0, value=1.1, step=0.1,
+                            1.0,
+                            10.0,
+                            value=1.1,
+                            step=0.1,
+                            label=" ",
+                        )
+
+                        t2i_negative = gr.Textbox(
                             label=" ",
+                            lines=2,
+                        )
+
+                        t2i_seed = gr.Slider(
+                            0,
+                            MAX_SEED,
+                            value=0,
+                            step=1,
+                            label=" ",
+                        )
+
+                        t2i_random = gr.Checkbox(
+                            label=" ",
+                            value=True,
                         )
-                        t2i_negative = gr.Textbox(label=" ", lines=2)
-                        t2i_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label=" ")
-                        t2i_random = gr.Checkbox(label=" ", value=True)
 
                 with gr.Column(scale=1):
-                    t2i_image = gr.Image(label=" ", type="pil", format="png", height=375)
-                    t2i_file = gr.File(label=" ", visible=False)
+                    t2i_image = gr.Image(
+                        label=" ",
+                        type="pil",
+                        format="png",
+                        height=375,
+                    )
+
+                    t2i_file = gr.File(
+                        label=" ",
+                        visible=False,
+                    )
+
                     with gr.Row():
-                        t2i_used_seed = gr.Number(label=" ", interactive=False, visible=False)
-                        t2i_timing = gr.Textbox(label=" ", interactive=False, visible=False)
+                        t2i_used_seed = gr.Number(
+                            label=" ",
+                            interactive=False,
+                            visible=False,
+                        )
+
+                        t2i_timing = gr.Textbox(
+                            label=" ",
+                            interactive=False,
+                            visible=False,
+                        )
 
             gr.Examples(
                 examples=[
-                    ['A neon shop sign that reads "QWEN IMAGE 2.1", rainy night, reflections on wet pavement', False, "16:9 — 1376x768"],
-                    ["A capybara wearing a wizard hat, reading a book by candlelight, oil painting", False, "1:1 — 1024x1024"],
-                    ["A cute cartoon dragon sticker", True, "1:1 — 1024x1024"],
-                    ["A glossy 3D icon of a paper plane, soft studio lighting", True, "1:1 — 1024x1024"],
-                    ["A handwritten chalkboard menu in a Parisian cafe, legible French text, warm morning light", False, "3:4 — 896x1200"],
+                    [
+                        'A neon shop sign that reads "QWEN IMAGE 2.1", '
+                        "rainy night, reflections on wet pavement",
+                        False,
+                        "16:9 — 1376x768",
+                    ],
+                    [
+                        "A capybara wearing a wizard hat, reading a book by "
+                        "candlelight, oil painting",
+                        False,
+                        "1:1 — 1024x1024",
+                    ],
+                    [
+                        "A cute cartoon dragon sticker",
+                        True,
+                        "1:1 — 1024x1024",
+                    ],
+                    [
+                        "A glossy 3D icon of a paper plane, soft studio lighting",
+                        True,
+                        "1:1 — 1024x1024",
+                    ],
+                    [
+                        "A handwritten chalkboard menu in a Parisian cafe, "
+                        "legible French text, warm morning light",
+                        False,
+                        "3:4 — 896x1200",
+                    ],
+                ],
+                inputs=[
+                    t2i_prompt,
+                    t2i_transparent,
+                    t2i_size,
+                ],
+                outputs=[
+                    t2i_image,
+                    t2i_file,
+                    t2i_used_seed,
+                    t2i_timing,
                 ],
-                inputs=[t2i_prompt, t2i_transparent, t2i_size],
-                outputs=[t2i_image, t2i_file, t2i_used_seed, t2i_timing],
                 fn=generate,
                 cache_examples=True,
                 cache_mode="lazy",
@@ -305,6 +541,7 @@ with gr.Blocks(title=" ") as demo:
                         label=" ",
                         lines=6,
                     )
+
                     edit_refs = gr.Gallery(
                         label=" ",
                         type="filepath",
@@ -312,47 +549,145 @@ with gr.Blocks(title=" ") as demo:
                         height=344,
                         interactive=True,
                     )
-                    edit_run = gr.Button(" ", variant="primary")
+
+                    edit_size = gr.Radio(
+                        label="Output size",
+                        choices=list(EDIT_SIZE_PRESETS),
+                        value=DEFAULT_EDIT_SIZE,
+                    )
+
+                    edit_run = gr.Button(
+                        " ",
+                        variant="primary",
+                    )
 
                     with gr.Accordion(" ", open=False):
-                        edit_steps = gr.Slider(10, 50, value=25, step=1, label=" ")
+                        edit_steps = gr.Slider(
+                            10,
+                            50,
+                            value=25,
+                            step=1,
+                            label=" ",
+                        )
+
                         edit_cfg = gr.Slider(
-                            1.0, 10.0, value=1.1, step=0.1,
+                            1.0,
+                            10.0,
+                            value=1.1,
+                            step=0.1,
                             label=" ",
                         )
-                        edit_negative = gr.Textbox(label=" ", lines=2)
-                        edit_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label=" ")
-                        edit_random = gr.Checkbox(label=" ", value=True)
+
+                        edit_negative = gr.Textbox(
+                            label=" ",
+                            lines=2,
+                        )
+
+                        edit_seed = gr.Slider(
+                            0,
+                            MAX_SEED,
+                            value=0,
+                            step=1,
+                            label=" ",
+                        )
+
+                        edit_random = gr.Checkbox(
+                            label=" ",
+                            value=True,
+                        )
 
                 with gr.Column(scale=1):
-                    edit_image = gr.Image(label=" ", type="pil", format="png", height=344)
-                    edit_file = gr.File(label=" ", visible=False)
+                    edit_image = gr.Image(
+                        label=" ",
+                        type="pil",
+                        format="png",
+                        height=344,
+                    )
+
+                    edit_file = gr.File(
+                        label=" ",
+                        visible=False,
+                    )
+
                     with gr.Row():
-                        edit_used_seed = gr.Number(label=" ", interactive=False, visible=False)
-                        edit_timing = gr.Textbox(label=" ", interactive=False, visible=False)
+                        edit_used_seed = gr.Number(
+                            label=" ",
+                            interactive=False,
+                            visible=False,
+                        )
 
-            gr.Markdown(
-                " "
-            )
+                        edit_timing = gr.Textbox(
+                            label=" ",
+                            interactive=False,
+                            visible=False,
+                        )
+
+            gr.Markdown(" ")
 
     t2i_run.click(
         fn=generate,
-        inputs=[t2i_prompt, t2i_transparent, t2i_size, t2i_steps, t2i_cfg, t2i_negative, t2i_seed, t2i_random],
-        outputs=[t2i_image, t2i_file, t2i_used_seed, t2i_timing],
+        inputs=[
+            t2i_prompt,
+            t2i_transparent,
+            t2i_size,
+            t2i_steps,
+            t2i_cfg,
+            t2i_negative,
+            t2i_seed,
+            t2i_random,
+        ],
+        outputs=[
+            t2i_image,
+            t2i_file,
+            t2i_used_seed,
+            t2i_timing,
+        ],
         api_name="generate",
     )
+
     t2i_prompt.submit(
         fn=generate,
-        inputs=[t2i_prompt, t2i_transparent, t2i_size, t2i_steps, t2i_cfg, t2i_negative, t2i_seed, t2i_random],
-        outputs=[t2i_image, t2i_file, t2i_used_seed, t2i_timing],
+        inputs=[
+            t2i_prompt,
+            t2i_transparent,
+            t2i_size,
+            t2i_steps,
+            t2i_cfg,
+            t2i_negative,
+            t2i_seed,
+            t2i_random,
+        ],
+        outputs=[
+            t2i_image,
+            t2i_file,
+            t2i_used_seed,
+            t2i_timing,
+        ],
         api_name=False,
     )
+
     edit_run.click(
         fn=edit,
-        inputs=[edit_instruction, edit_refs, edit_steps, edit_cfg, edit_negative, edit_seed, edit_random],
-        outputs=[edit_image, edit_file, edit_used_seed, edit_timing],
+        inputs=[
+            edit_instruction,
+            edit_refs,
+            edit_size,
+            edit_steps,
+            edit_cfg,
+            edit_negative,
+            edit_seed,
+            edit_random,
+        ],
+        outputs=[
+            edit_image,
+            edit_file,
+            edit_used_seed,
+            edit_timing,
+        ],
         api_name="edit",
     )
 
+
 if __name__ == "__main__":
     demo.launch(mcp_server=False)
+