Update app.py

app.py

@@ -1,9 +1,10 @@
 import os, json, uuid, re
 from datetime import datetime
 import gradio as gr
-import spaces  # ZeroGPU decorator
+import spaces
 import torch
 from PIL import Image
+import pandas as pd
 
 # =========================
 # Storage helpers
@@ -30,7 +31,7 @@ def save_project(proj):
 def load_project_file(file_obj):
     with open(file_obj.name, "r") as f:
         proj = json.load(f)
-    project_dir(proj["meta"]["id"])  # ensure dirs
+    project_dir(proj["meta"]["id"])
     return proj
 
 def ensure_project(p, suggested_name="Project"):
@@ -40,14 +41,14 @@ def ensure_project(p, suggested_name="Project"):
     name = f"{suggested_name}-{pid[:4]}"
     proj = {
         "meta": {"id": pid, "name": name, "created": now_iso(), "updated": now_iso()},
-        "shots": [],    # each: id,title,description,duration,fps,steps,seed,negative, image_path?(on approval)
+        "shots": [],    # each shot: id,title,description,duration,fps,steps,seed,negative,image_path?
         "clips": []
     }
     save_project(proj)
     return proj
 
 # =========================
-# LLM (ZeroGPU) — Storyboard generator (robust, two-pass + empty fallback)
+# LLM (ZeroGPU) — Storyboard generator (robust)
 # =========================
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
@@ -119,6 +120,7 @@ def _apply_chat(tok, system_msg: str, user_msg: str) -> str:
     return system_msg + "\n\n" + user_msg
 
 def _generate_text(model, tok, prompt_text: str) -> str:
+    """Decode only the continuation (avoid prompt echo)."""
     inputs = tok(prompt_text, return_tensors="pt")
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
     eos_id = tok.eos_token_id or tok.pad_token_id
@@ -132,7 +134,6 @@ def _generate_text(model, tok, prompt_text: str) -> str:
         eos_token_id=eos_id,
         pad_token_id=eos_id,
     )
-    # decode only continuation
     prompt_len = inputs["input_ids"].shape[1]
     continuation_ids = gen[0][prompt_len:]
     text = tok.decode(continuation_ids, skip_special_tokens=True).strip()
@@ -146,7 +147,6 @@ def _extract_json_array(text: str) -> str:
         inner = m.group(1).strip()
         if inner:
             return inner
-    # Fallback: first balanced array
     start = text.find("[")
     if start == -1:
         return ""
@@ -173,7 +173,7 @@ def _normalize_shots(shots_raw, default_fps: int, default_len: int):
             "steps": int(s.get("steps", 30)),
             "seed": s.get("seed", None),
             "negative": s.get("negative", ""),
-            "image_path": s.get("image_path", None)  # will be set after approval
+            "image_path": s.get("image_path", None)
         })
     return norm
 
@@ -182,13 +182,13 @@ def generate_storyboard_with_llm(user_prompt: str, n_shots: int, default_fps: in
     model, tok = _lazy_model_tok()
     system = "You are a film previsualization assistant. Output must be valid JSON."
 
-    # PASS 1
+    # Pass 1
     p1 = _apply_chat(tok, system + " Return ONLY JSON inside <JSON> tags.",
                      _prompt_with_tags(user_prompt, n_shots, default_fps, default_len))
     out1 = _generate_text(model, tok, p1)
     json_text = _extract_json_array(out1)
 
-    # PASS 2 fallback
+    # Pass 2
     if not json_text:
         p2 = _apply_chat(tok, system + " Reply ONLY with a JSON array.",
                          _prompt_minimal(user_prompt, n_shots, default_fps, default_len))
@@ -199,7 +199,7 @@ def generate_storyboard_with_llm(user_prompt: str, n_shots: int, default_fps: in
             if start != -1 and end != -1 and end > start:
                 json_text = out2[start:end+1].strip()
 
-    # EMPTY FALLBACK: simple storyboard so UI never crashes
+    # Empty fallback
     if not json_text or not json_text.strip():
         fallback = []
         for i in range(1, int(n_shots) + 1):
@@ -234,20 +234,35 @@ _sd_t2i = None
 _sd_i2i = None
 
 def _lazy_sd_pipes():
+    """Load SD once, disable safety checker to avoid offload_state_dict issues; reuse modules for img2img."""
     global _sd_t2i, _sd_i2i
     if _sd_t2i is not None and _sd_i2i is not None:
         return _sd_t2i, _sd_i2i
+
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
     _sd_t2i = StableDiffusionPipeline.from_pretrained(
-        SD_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-    )
-    _sd_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
-        SD_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        SD_MODEL,
+        torch_dtype=dtype,
+        safety_checker=None,
+        feature_extractor=None,
+        use_safetensors=True
     )
     if torch.cuda.is_available():
         _sd_t2i = _sd_t2i.to("cuda")
+
+    _sd_i2i = StableDiffusionImg2ImgPipeline(
+        vae=_sd_t2i.vae,
+        text_encoder=_sd_t2i.text_encoder,
+        tokenizer=_sd_t2i.tokenizer,
+        unet=_sd_t2i.unet,
+        scheduler=_sd_t2i.scheduler,
+        safety_checker=None,
+        feature_extractor=None
+    )
+    if torch.cuda.is_available():
         _sd_i2i = _sd_i2i.to("cuda")
-    _sd_t2i.safety_checker = None
-    _sd_i2i.safety_checker = None
+
     return _sd_t2i, _sd_i2i
 
 def _save_keyframe(pid: str, shot_id: int, img: Image.Image) -> str:
@@ -266,9 +281,8 @@ def generate_keyframe_image(
 ):
     """
     Generate image for shots[shot_idx].
-    - If shot_idx == 0: text2img
-    - Else: img2img with previous shot's approved image_path as init image
-    Uses edited fields in shots: description, negative, steps, seed.
+    - shot 0: text2img
+    - shot k>0: img2img using previous approved image as conditioning (if available)
     """
     t2i, i2i = _lazy_sd_pipes()
     shot = shots[shot_idx]
@@ -276,21 +290,31 @@ def generate_keyframe_image(
     negative = shot.get("negative") or ""
     steps = int(shot.get("steps", 30))
     seed = shot.get("seed", None)
+
     gen = torch.Generator("cuda" if torch.cuda.is_available() else "cpu")
     if isinstance(seed, int):
         gen = gen.manual_seed(seed)
 
     if shot_idx == 0 or not shots[shot_idx - 1].get("image_path"):
-        # text2img
-        out = t2i(prompt=prompt, negative_prompt=negative, guidance_scale=guidance_scale,
-                  num_inference_steps=steps, generator=gen).images[0]
+        out = t2i(
+            prompt=prompt,
+            negative_prompt=negative,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps,
+            generator=gen
+        ).images[0]
     else:
-        # img2img: previous approved keyframe as conditioning
         prev_path = shots[shot_idx - 1]["image_path"]
         init_image = Image.open(prev_path).convert("RGB")
-        out = i2i(prompt=prompt, negative_prompt=negative, image=init_image,
-                  guidance_scale=guidance_scale, strength=strength,
-                  num_inference_steps=steps, generator=gen).images[0]
+        out = i2i(
+            prompt=prompt,
+            negative_prompt=negative,
+            image=init_image,
+            guidance_scale=guidance_scale,
+            strength=strength,
+            num_inference_steps=steps,
+            generator=gen
+        ).images[0]
 
     saved_path = _save_keyframe(pid, int(shot["id"]), out)
     return saved_path
@@ -298,23 +322,18 @@ def generate_keyframe_image(
 # =========================
 # Shots <-> Dataframe utils
 # =========================
-import pandas as pd
-
 SHOT_COLUMNS = ["id", "title", "description", "duration", "fps", "steps", "seed", "negative", "image_path"]
 
 def shots_to_df(shots: list) -> pd.DataFrame:
-    rows = []
-    for s in shots:
-        rows.append({k: s.get(k, None) for k in SHOT_COLUMNS})
-    df = pd.DataFrame(rows, columns=SHOT_COLUMNS)
-    return df
+    rows = [{k: s.get(k, None) for k in SHOT_COLUMNS} for s in shots]
+    return pd.DataFrame(rows, columns=SHOT_COLUMNS)
 
 def df_to_shots(df: pd.DataFrame) -> list:
     out = []
     for _, row in df.iterrows():
         out.append({
             "id": int(row["id"]),
-            "title": row["title"] or f"Shot {int(row['id'])}",
+            "title": (row["title"] or f"Shot {int(row['id'])}"),
             "description": row["description"] or "",
             "duration": int(row["duration"]) if pd.notna(row["duration"]) else 4,
             "fps": int(row["fps"]) if pd.notna(row["fps"]) else 24,
@@ -323,22 +342,20 @@ def df_to_shots(df: pd.DataFrame) -> list:
             "negative": row["negative"] or "",
             "image_path": row["image_path"] if pd.notna(row["image_path"]) else None
         })
-    # keep sorted by id
-    out = sorted(out, key=lambda x: x["id"])
-    return out
+    return sorted(out, key=lambda x: x["id"])
 
 # =========================
 # Gradio UI
 # =========================
 with gr.Blocks() as demo:
     gr.Markdown("# 🎬 Storyboard → Keyframes → Videos → Export")
-    gr.Markdown("**Step 3**: Edit storyboard, then generate keyframes. Shot 2..N use the previous approved image as reference (img2img).")
+    gr.Markdown("**Edit storyboard prompts**, then generate keyframes. Each next shot uses the **previous approved image** as reference.")
 
-    # Global state
-    project = gr.State(None)        # dict with meta/shots/clips
-    current_idx = gr.State(0)       # index of current shot in Keyframes tab
+    # State
+    project = gr.State(None)
+    current_idx = gr.State(0)
 
-    # Header row
+    # Header
     with gr.Row():
         with gr.Column(scale=2):
             proj_name = gr.Textbox(label="Project name", placeholder="e.g., Desert Chase")
@@ -349,6 +366,7 @@ with gr.Blocks() as demo:
         with gr.Column(scale=1):
             load_file = gr.File(label="Load Project (project.json)", file_count="single", type="filepath")
             load_btn = gr.Button("Load")
+    sb_status = gr.Markdown("")
 
     # Tabs
     with gr.Tabs():
@@ -360,19 +378,21 @@ with gr.Blocks() as demo:
                 sb_default_fps  = gr.Slider(8, 60, value=24, step=1, label="Default FPS")
                 sb_default_len  = gr.Slider(1, 12, value=4, step=1, label="Default seconds per shot")
             propose_btn = gr.Button("Propose Storyboard (LLM on ZeroGPU)")
-            shots_df    = gr.Dataframe(headers=SHOT_COLUMNS, datatype=["number","str","str","number","number","number","number","str","str"], row_count=(1,"dynamic"), col_count=len(SHOT_COLUMNS), label="Edit shots below", wrap=True)
-            save_edits_btn = gr.Button("Save Edits ✓", variant="primary")
+            shots_df    = gr.Dataframe(
+                headers=SHOT_COLUMNS,
+                datatype=["number","str","str","number","number","number","number","str","str"],
+                row_count=(1,"dynamic"), col_count=len(SHOT_COLUMNS),
+                label="Edit shots below (prompts & params)", wrap=True
+            )
+            save_edits_btn = gr.Button("Save Edits ✓", variant="primary", interactive=False)
             to_keyframes_btn = gr.Button("Start Keyframes →", variant="secondary")
-            sb_status    = gr.Markdown("")
 
         with gr.Tab("Keyframes"):
             gr.Markdown("### 2) Keyframes")
+            shot_info_md = gr.Markdown("")
+            prompt_box   = gr.Textbox(label="Shot description (editable before generating)", lines=4)
             with gr.Row():
-                shot_info_md = gr.Markdown("")
-            with gr.Row():
-                prompt_box = gr.Textbox(label="Shot description (editable before generating)", lines=4)
-            with gr.Row():
-                gen_btn = gr.Button("Generate / Regenerate (uses previous approved image if available)", variant="primary")
+                gen_btn = gr.Button("Generate / Regenerate", variant="primary")
                 approve_next_btn = gr.Button("Approve & Next →", variant="secondary")
             with gr.Row():
                 prev_img = gr.Image(label="Previous approved image (conditioning)", type="filepath")
@@ -387,7 +407,7 @@ with gr.Blocks() as demo:
             gr.Markdown("### 4) Export (coming next)")
             export_info = gr.Markdown("Nothing to export yet.")
 
-    # -------- Handlers --------
+    # ---------- Handlers ----------
     def on_new(name):
         p = ensure_project(None, suggested_name=(name or "Project"))
         return p, gr.update(value=f"**New project created** `{p['meta']['name']}` (id: `{p['meta']['id']}`)")
@@ -403,17 +423,23 @@ with gr.Blocks() as demo:
         p["shots"] = shots
         p["meta"]["updated"] = now_iso()
         save_project(p)
-        return p, shots_to_df(shots), gr.update(value="Storyboard generated (editable).")
+        # Enable Save Edits after storyboard exists
+        return p, shots_to_df(shots), gr.update(value="Storyboard generated (editable)."), gr.update(interactive=True)
 
     propose_btn.click(
         on_propose,
         inputs=[project, sb_prompt, sb_target_shots, sb_default_fps, sb_default_len],
-        outputs=[project, shots_df, sb_status]
+        outputs=[project, shots_df, sb_status, save_edits_btn]
     )
 
-    def on_save_edits(p, df):
+    # Defensive save handler (works even if user clicks too early)
+    def on_save_edits(*args):
+        p = args[0] if len(args) > 0 else None
+        df = args[1] if len(args) > 1 else None
         if p is None:
-            raise gr.Error("No project in memory.")
+            raise gr.Error("No project in memory. Click New Project, then generate a storyboard.")
+        if df is None:
+            raise gr.Error("No storyboard table to save. Generate a storyboard first, then edit it.")
         shots = df_to_shots(df)
         p = dict(p)
         p["shots"] = shots
@@ -439,8 +465,7 @@ with gr.Blocks() as demo:
         if p is None: raise gr.Error("No project.")
         shots = p["shots"]
         if idx < 0 or idx >= len(shots): raise gr.Error("Invalid shot index.")
-        # Allow in-place prompt tweak before generation
-        shots[idx]["description"] = current_prompt
+        shots[idx]["description"] = current_prompt  # allow tweaking before generation
         prev_path = shots[idx-1]["image_path"] if idx > 0 else None
         img_path = generate_keyframe_image(p["meta"]["id"], int(idx), shots)
         return img_path, (prev_path or None), gr.update(value=f"Generated candidate for shot {shots[idx]['id']}.")
@@ -453,21 +478,20 @@ with gr.Blocks() as demo:
         i = int(idx)
         if i < 0 or i >= len(shots): raise gr.Error("Invalid shot index.")
         if not latest_img_path: raise gr.Error("Generate an image first.")
-        # commit prompt and image path
+        # commit
         shots[i]["description"] = current_prompt
         shots[i]["image_path"] = latest_img_path
         p["shots"] = shots
         p["meta"]["updated"] = now_iso()
         save_project(p)
 
-        # Move to next
+        # next
         if i + 1 < len(shots):
             ni = i + 1
             info = f"**Shot {shots[ni]['id']} — {shots[ni]['title']}**  \nDuration: {shots[ni]['duration']}s @ {shots[ni]['fps']} fps"
             prev_path = shots[ni-1]["image_path"]
             return p, ni, gr.update(value=info), gr.update(value=shots[ni]["description"]), gr.update(value=prev_path), gr.update(value=None), gr.update(value=f"Approved shot {shots[i]['id']}. On to shot {shots[ni]['id']}.")
         else:
-            # finished all keyframes
             return p, i, gr.update(value="**All keyframes approved.** Proceed to Videos tab."), gr.update(value=""), gr.update(value=shots[i]["image_path"]), gr.update(value=None), gr.update(value="All shots approved ✅")
 
     approve_next_btn.click(on_approve_next, inputs=[project, current_idx, prompt_box, out_img], outputs=[project, current_idx, shot_info_md, prompt_box, prev_img, out_img, kf_status])