Add Extend prompt + custom start image controls (app.py)

app.py

@@ -33,6 +33,9 @@ DEFAULT_NEGATIVE = (
     "静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
 )
 
+EXTEND_START_LAST_FRAME = "Last frame from video"
+EXTEND_START_UPLOAD = "Upload custom image"
+
 WORK = Path(tempfile.gettempdir()) / "wan22_extend"
 WORK.mkdir(parents=True, exist_ok=True)
 
@@ -261,6 +264,41 @@ def estimate_total(seg_duration: float, num_segments: int) -> float:
     return float(seg_duration) * int(num_segments)
 
 
+def resolve_extend_prompt(extend_prompt: str | None, generate_prompt: str | None, state: dict) -> str:
+    """Prefer dedicated Extend prompt; else Generate prompt / saved state / default."""
+    if extend_prompt and str(extend_prompt).strip():
+        return str(extend_prompt).strip()
+    if generate_prompt and str(generate_prompt).strip():
+        return str(generate_prompt).strip()
+    return state.get("prompt") or DEFAULT_PROMPT
+
+
+def resolve_extend_start_image(
+    extend_start: str,
+    custom_image,
+    state: dict,
+    sess: Path,
+) -> Path:
+    """Pick I2V start image for an Extend segment."""
+    mode = (extend_start or EXTEND_START_LAST_FRAME).strip()
+    if mode == EXTEND_START_UPLOAD:
+        if custom_image is None:
+            raise gr.Error("Upload a custom extend start image, or switch to “Last frame from video”.")
+        saved = save_image(custom_image, sess / f"extend_custom_{uuid.uuid4().hex[:8]}")
+        if saved is None:
+            raise gr.Error("Could not read the custom extend start image.")
+        return saved
+
+    # Default: last frame from current video
+    last = Path(state["last_frame"]) if state.get("last_frame") else None
+    if last is None or not last.exists():
+        if not state.get("video"):
+            raise gr.Error("Generate a clip first, then use Extend.")
+        last = extract_last_frame(state["video"], sess / "last_frame")
+        state["last_frame"] = str(last)
+    return last
+
+
 def do_generate(
     image,
     prompt,
@@ -325,6 +363,9 @@ def do_generate(
 
 def do_extend(
     prompt,
+    extend_prompt,
+    extend_start,
+    custom_extend_image,
     seg_duration,
     steps,
     negative,
@@ -336,22 +377,20 @@ def do_extend(
     state: dict,
     progress=gr.Progress(track_tqdm=False),
 ):
-    if not state or not state.get("video") or not state.get("last_frame"):
+    if not state or not state.get("video"):
         raise gr.Error("Generate a clip first, then use Extend.")
     if len(state.get("segments", [])) >= 6:
         raise gr.Error("Segment cap reached (6). Reset and start a new chain.")
 
     sess = Path(state["dir"])
-    last = Path(state["last_frame"])
-    if not last.exists():
-        last = extract_last_frame(state["video"], sess / "last_frame")
-        state["last_frame"] = str(last)
+    start_img = resolve_extend_start_image(extend_start, custom_extend_image, state, sess)
+    use_prompt = resolve_extend_prompt(extend_prompt, prompt, state)
 
-    use_prompt = prompt or state.get("prompt") or DEFAULT_PROMPT
-    progress(0.05, desc="Extending from last frame…")
+    mode_label = (extend_start or EXTEND_START_LAST_FRAME).strip()
+    progress(0.05, desc=f"Extending ({mode_label})…")
     try:
         clip = call_wan_i2v(
-            last,
+            start_img,
             use_prompt,
             float(seg_duration),
             int(steps),
@@ -373,6 +412,8 @@ def do_extend(
     except Exception as e:
         raise gr.Error(f"Concat failed (ffmpeg): {e}") from e
 
+    # Always refresh last_frame from the new concatenated video so later
+    # “Last frame from video” extends stay correct.
     last = extract_last_frame(out, sess / "last_frame")
     dur = video_duration_seconds(out)
     state["segments"] = [str(p) for p in segs]
@@ -392,6 +433,9 @@ def do_extend(
 def do_auto_extend(
     target_seconds,
     prompt,
+    extend_prompt,
+    extend_start,
+    custom_extend_image,
     seg_duration,
     steps,
     negative,
@@ -414,6 +458,7 @@ def do_auto_extend(
         state.get("last_frame"),
         state,
     )
+    first = True
     while len(state.get("segments", [])) < max_segs:
         dur = video_duration_seconds(state["video"])
         if dur >= target:
@@ -422,8 +467,15 @@ def do_auto_extend(
             len(state["segments"]) / max_segs,
             desc=f"Auto-extend: {dur:.1f}s / {target:.0f}s…",
         )
+        # Custom image only applies to the first auto-extend segment;
+        # later segments chain from the updated video last frame.
+        this_start = extend_start if first else EXTEND_START_LAST_FRAME
+        this_custom = custom_extend_image if first else None
         video, download, status, last_preview, state = do_extend(
             prompt,
+            extend_prompt,
+            this_start,
+            this_custom,
             seg_duration,
             steps,
             negative,
@@ -435,6 +487,7 @@ def do_auto_extend(
             state,
             progress,
         )
+        first = False
         time.sleep(0.3)
     return video, download, status, last_preview, state
 
@@ -453,6 +506,11 @@ def update_estimate(seg_duration, num_planned):
     return f"Planned length ≈ **{total:.1f}s** ({num_planned} × {seg_duration:.1f}s segments). Cap: 6 segments."
 
 
+def toggle_custom_extend_image(mode: str):
+    visible = (mode or "").strip() == EXTEND_START_UPLOAD
+    return gr.update(visible=visible)
+
+
 CSS = """
 .big-btn button { font-size: 1.15rem !important; min-height: 3rem !important; }
 #status-md { font-size: 1rem; }
@@ -465,7 +523,7 @@ def build_ui():
         gr.Markdown(
             """
 # Wan 2.2 Image → Video + Extend
-Upload a still, generate a short Wan 2.2 clip on **free ZeroGPU**, then **Extend** from the last frame to chain ~10–20s videos.
+Upload a still, generate a short Wan 2.2 clip on **free ZeroGPU**, then **Extend** from the last frame (or a custom start image) to chain ~10–20s videos.
 
 ZeroGPU is free but **queued and quota-limited** — each segment often takes 1–3 minutes when busy.
             """
@@ -476,11 +534,34 @@ ZeroGPU is free but **queued and quota-limited** — each segment often takes 1
             with gr.Column(scale=1):
                 image = gr.Image(label="Still image", type="pil", height=320)
                 prompt = gr.Textbox(
-                    label="Prompt",
+                    label="Generate prompt",
                     value=DEFAULT_PROMPT,
                     lines=3,
                     placeholder="Describe the motion you want…",
                 )
+
+                with gr.Accordion("Extend options", open=True):
+                    extend_prompt = gr.Textbox(
+                        label="Extend prompt",
+                        value="",
+                        lines=3,
+                        placeholder="Leave blank to reuse Generate prompt",
+                        info="Optional. Used only for Extend / Auto-extend.",
+                    )
+                    extend_start = gr.Radio(
+                        choices=[EXTEND_START_LAST_FRAME, EXTEND_START_UPLOAD],
+                        value=EXTEND_START_LAST_FRAME,
+                        label="Extend start",
+                        info="Default keeps current behavior (auto last frame).",
+                    )
+                    custom_extend_image = gr.Image(
+                        label="Custom extend start image",
+                        type="pil",
+                        height=220,
+                        visible=False,
+                        info="Replaces the auto last frame as the I2V start for the next segment.",
+                    )
+
                 with gr.Row():
                     seg_duration = gr.Slider(
                         2.0,
@@ -543,6 +624,7 @@ ZeroGPU is free but **queued and quota-limited** — each segment often takes 1
 
         seg_duration.change(update_estimate, [seg_duration, num_planned], estimate)
         num_planned.change(update_estimate, [seg_duration, num_planned], estimate)
+        extend_start.change(toggle_custom_extend_image, [extend_start], [custom_extend_image])
 
         gen_inputs = [
             image,
@@ -559,6 +641,9 @@ ZeroGPU is free but **queued and quota-limited** — each segment often takes 1
         ]
         ext_inputs = [
             prompt,
+            extend_prompt,
+            extend_start,
+            custom_extend_image,
             seg_duration,
             steps,
             negative,
@@ -583,8 +668,8 @@ ZeroGPU is free but **queued and quota-limited** — each segment often takes 1
         gr.Markdown(
             f"""
 ### How Extend works
-1. **Generate** — still image + prompt → short Wan 2.2 I2V clip (~{3.5}s default).
-2. **Extend** — take the **last frame**, run I2V again, **concatenate** onto the previous video.
+1. **Generate** — still image + Generate prompt → short Wan 2.2 I2V clip (~{3.5}s default).
+2. **Extend** — start from the **last frame** (or a **custom uploaded image**), optional separate **Extend prompt**, run I2V again, **concatenate**.
 3. Repeat until you hit ~10–20s (or use **Auto-extend**). Soft cap: **6 segments**.
 
 ### Limits