Update app.py

app.py

@@ -1,15 +1,18 @@
 import os, json, uuid
 from datetime import datetime
 import gradio as gr
+import spaces  # <<< required for ZeroGPU
 
-# -------- storage helpers --------
+# =========================
+# Storage helpers
+# =========================
 ROOT = "outputs"
 os.makedirs(ROOT, exist_ok=True)
 
-def now_iso(): return datetime.utcnow().replace(microsecond=0).isoformat()+"Z"
+def now_iso(): return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
 def new_id(): return uuid.uuid4().hex[:8]
 
-def project_dir(pid): 
+def project_dir(pid):
     path = os.path.join(ROOT, pid)
     os.makedirs(path, exist_ok=True)
     os.makedirs(os.path.join(path, "keyframes"), exist_ok=True)
@@ -25,39 +28,111 @@ def save_project(proj):
 def load_project_file(file_obj):
     with open(file_obj.name, "r") as f:
         proj = json.load(f)
-    # ensure directories exist
-    project_dir(proj["meta"]["id"])
+    project_dir(proj["meta"]["id"])  # ensure dirs
     return proj
 
-# -------- model-free stubs for Step 1 --------
-def stub_storyboard_from_prompt(prompt, target_shots, default_fps, default_len):
-    """Return a minimal storyboard (no LLM yet)."""
-    shots = []
-    n = target_shots or 3
-    for i in range(1, n+1):
-        shots.append({
-            "id": i,
-            "title": f"Shot {i}",
-            "description": f"Placeholder description for shot {i} derived from: {prompt[:80]}",
-            "duration": default_len,
-            "fps": default_fps,
-            "video_length": default_len,
-            "steps": 30,
-            "seed": None,
-            "negative": "",
+# =========================
+# LLM (ZeroGPU) — Storyboard generator
+# =========================
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+STORYBOARD_MODEL = os.getenv("STORYBOARD_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
+HF_TASK_MAX_TOKENS = int(os.getenv("HF_TASK_MAX_TOKENS", "900"))  # keep tidy for JSON
+_pipe = None  # lazy-loaded global
+
+def _lazy_pipe():
+    global _pipe
+    if _pipe is not None:
+        return _pipe
+    tok = AutoTokenizer.from_pretrained(STORYBOARD_MODEL, trust_remote_code=True)
+    mdl = AutoModelForCausalLM.from_pretrained(
+        STORYBOARD_MODEL,
+        device_map="auto",
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+    _pipe = pipeline(
+        "text-generation",
+        model=mdl,
+        tokenizer=tok,
+        max_new_tokens=HF_TASK_MAX_TOKENS,
+        do_sample=False,              # deterministic JSON
+        temperature=0.0,
+        repetition_penalty=1.05,
+    )
+    return _pipe
+
+def _storyboard_prompt(user_prompt: str, n_shots: int, default_fps: int, default_len: int) -> str:
+    return f"""
+You are a film previsualization assistant. Return ONLY valid JSON (no explanations).
+Create a storyboard of {n_shots} numbered shots for the following idea:
+
+\"\"\"{user_prompt}\"\"\"
+
+Return an array of objects with this exact schema and default values:
+[
+  {{
+    "id": 1,
+    "title": "Short title",
+    "description": "A visual description suitable for keyframe generation",
+    "duration": {default_len},
+    "fps": {default_fps},
+    "video_length": {default_len},
+    "steps": 30,
+    "seed": null,
+    "negative": ""
+  }}
+]
+
+Rules:
+- IDs must start at 1 and increment by 1.
+- Use simple ASCII only. No trailing commas.
+- Output must be valid JSON parseable by Python's json.loads.
+""".strip()
+
+@spaces.GPU(duration=180)  # <<< ZeroGPU entrypoint: triggers pooled GPU allocation
+def generate_storyboard_with_llm(user_prompt: str, n_shots: int, default_fps: int, default_len: int):
+    pipe = _lazy_pipe()
+    prompt = _storyboard_prompt(user_prompt, n_shots, default_fps, default_len)
+    out = pipe(prompt)[0]["generated_text"]
+
+    # Extract the JSON array
+    start = out.find("[")
+    end = out.rfind("]")
+    if start == -1 or end == -1 or end <= start:
+        raise ValueError("LLM did not return valid JSON.")
+    text = out[start:end+1]
+    shots = json.loads(text)
+
+    # Normalize & enforce required fields
+    norm = []
+    for i, s in enumerate(shots, start=1):
+        norm.append({
+            "id": int(s.get("id", i)),
+            "title": s.get("title", f"Shot {i}"),
+            "description": s.get("description", ""),
+            "duration": int(s.get("duration", default_len)),
+            "fps": int(s.get("fps", default_fps)),
+            "video_length": int(s.get("video_length", default_len)),
+            "steps": int(s.get("steps", 30)),
+            "seed": s.get("seed", None),
+            "negative": s.get("negative", ""),
             "keyframe_path": None
         })
-    return shots
+    return norm
 
-# -------- Gradio app --------
+# =========================
+# Gradio UI
+# =========================
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎬 Storyboard → Keyframes → Videos → Export (Scaffold)")
-    gr.Markdown("This is the skeleton UI. In the next steps we’ll plug in an open LLM, img2img keyframes, and your Modal video generator.")
+    gr.Markdown("# ��� Storyboard → Keyframes → Videos → Export")
+    gr.Markdown("**Step 2**: Real storyboard generation on **ZeroGPU**. Next steps will add keyframes (img2img) and your Modal videos.")
 
     # Global state
-    project = gr.State(None)  # dict with meta/shots/clips
+    project = gr.State(None)        # dict with meta/shots/clips
     current_tab = gr.State("Storyboard")
 
+    # Header row
     with gr.Row():
         with gr.Column(scale=2):
             proj_name = gr.Textbox(label="Project name", placeholder="e.g., Desert Chase")
@@ -69,69 +144,59 @@ with gr.Blocks() as demo:
             load_file = gr.File(label="Load Project (project.json)", file_count="single", type="filepath")
             load_btn = gr.Button("Load")
 
-    # ----- Tabs -----
-    with gr.Tabs() as tabs:
+    # Tabs
+    with gr.Tabs():
         with gr.Tab("Storyboard"):
             gr.Markdown("### 1) Storyboard")
-            sb_prompt = gr.Textbox(label="High-level prompt", lines=3, placeholder="Describe the story you want to create…")
+            sb_prompt = gr.Textbox(label="High-level prompt", lines=4, placeholder="Describe the story you want to create…")
             with gr.Row():
                 sb_target_shots = gr.Slider(1, 12, value=3, step=1, label="Target # of shots")
-                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 (stub)")
-            shots_json = gr.JSON(label="Storyboard (editable later)")
-            confirm_btn = gr.Button("Confirm Storyboard ✓", variant="primary")
-            sb_status = gr.Markdown("")
+                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_json   = gr.JSON(label="Storyboard JSON (editable in next step)")
+            confirm_btn  = gr.Button("Confirm Storyboard ✓", variant="primary")
+            sb_status    = gr.Markdown("")
 
         with gr.Tab("Keyframes"):
-            gr.Markdown("### 2) Keyframes")
-            gr.Markdown("In Step 2–3 we’ll add img2img generation and approvals here.")
+            gr.Markdown("### 2) Keyframes (coming next)")
             kf_table = gr.JSON(label="Shots (read-only for now)")
             to_videos_btn = gr.Button("Continue to Videos →", interactive=False)
 
         with gr.Tab("Videos"):
-            gr.Markdown("### 3) Videos (Transitions)")
-            gr.Markdown("In Step 3–4 we’ll add your Modal start/end-frame video generation here.")
+            gr.Markdown("### 3) Videos (coming next)")
             vd_table = gr.JSON(label="Planned clip edges (read-only for now)")
             to_export_btn = gr.Button("Continue to Export →", interactive=False)
 
         with gr.Tab("Export"):
-            gr.Markdown("### 4) Export")
-            gr.Markdown("Final concatenation & asset packaging will land here in Step 4.")
+            gr.Markdown("### 4) Export (coming next)")
             export_info = gr.Markdown("Nothing to export yet.")
 
-    # ----- Handlers -----
+    # -------- Handlers --------
     def on_new(name):
-        name = name.strip() or f"Project-{new_id()}"
+        name = (name or "").strip() or f"Project-{new_id()}"
         pid = new_id()
         p = {
-            "meta": {
-                "id": pid,
-                "name": name,
-                "created": now_iso(),
-                "updated": now_iso()
-            },
+            "meta": {"id": pid, "name": name, "created": now_iso(), "updated": now_iso()},
             "shots": [],
             "clips": []
         }
         save_project(p)
         return p, gr.update(value=f"**New project created** `{name}` (id: `{pid}`)")
 
-    new_btn.click(
-        on_new,
-        inputs=[proj_name],
-        outputs=[project, sb_status]
-    )
+    new_btn.click(on_new, inputs=[proj_name], outputs=[project, sb_status])
 
     def on_propose(p, prompt, target_shots, fps, vlen):
         if p is None:
             raise gr.Error("Create a project first (New Project).")
-        shots = stub_storyboard_from_prompt(prompt or "", int(target_shots), int(fps), int(vlen))
+        if not prompt or not str(prompt).strip():
+            raise gr.Error("Please enter a high-level prompt.")
+        shots = generate_storyboard_with_llm(str(prompt).strip(), int(target_shots), int(fps), int(vlen))
         p = dict(p)
         p["shots"] = shots
         p["meta"]["updated"] = now_iso()
         save_project(p)
-        return p, shots, gr.update(value="Storyboard proposed (stub). Edit will be available in Step 2.")
+        return p, shots, gr.update(value="Storyboard generated by LLM (ZeroGPU).")
 
     propose_btn.click(
         on_propose,
@@ -142,10 +207,11 @@ with gr.Blocks() as demo:
     def on_confirm(p):
         if p is None or not p.get("shots"):
             raise gr.Error("No storyboard yet.")
-        # Build planned clip edges (i->i+1)
         edges = []
         for i in range(len(p["shots"]) - 1):
-            edges.append({"from": p["shots"][i]["id"], "to": p["shots"][i+1]["id"], "prompt": f"Transition from shot {i+1} to {i+2}"})
+            a = p["shots"][i]["id"]
+            b = p["shots"][i+1]["id"]
+            edges.append({"from": a, "to": b, "prompt": f"Transition from shot {a} to {b}"})
         p = dict(p)
         p["clips"] = edges
         p["meta"]["updated"] = now_iso()
@@ -170,11 +236,7 @@ with gr.Blocks() as demo:
         path = save_project(p)
         return gr.update(value=f"Saved to `{path}`")
 
-    save_btn.click(
-        on_save,
-        inputs=[project],
-        outputs=[sb_status]
-    )
+    save_btn.click(on_save, inputs=[project], outputs=[sb_status])
 
     def on_load(file_obj):
         p = load_project_file(file_obj)
@@ -193,4 +255,5 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
+    # SSR is fine; you can set share=True if you want a public link automatically
     demo.launch()