app.py

import os, json, uuid
from datetime import datetime
import gradio as gr
import spaces  # <<< required for ZeroGPU

# =========================
# Storage helpers
# =========================
ROOT = "outputs"
os.makedirs(ROOT, exist_ok=True)

def now_iso(): return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
def new_id(): return uuid.uuid4().hex[:8]

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)
    os.makedirs(os.path.join(path, "clips"), exist_ok=True)
    return path

def save_project(proj):
    pid = proj["meta"]["id"]
    path = os.path.join(project_dir(pid), "project.json")
    with open(path, "w") as f: json.dump(proj, f, indent=2)
    return path

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
    return proj

# =========================
# 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 norm

# =========================
# Gradio UI
# =========================
with gr.Blocks() as demo:
    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
    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")
        with gr.Column(scale=1):
            new_btn = gr.Button("New Project", variant="primary")
        with gr.Column(scale=1):
            save_btn = gr.Button("Save Project")
        with gr.Column(scale=1):
            load_file = gr.File(label="Load Project (project.json)", file_count="single", type="filepath")
            load_btn = gr.Button("Load")

    # Tabs
    with gr.Tabs():
        with gr.Tab("Storyboard"):
            gr.Markdown("### 1) Storyboard")
            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 (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 (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 (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 (coming next)")
            export_info = gr.Markdown("Nothing to export yet.")

    # -------- Handlers --------
    def on_new(name):
        name = (name or "").strip() or f"Project-{new_id()}"
        pid = new_id()
        p = {
            "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])

    def on_propose(p, prompt, target_shots, fps, vlen):
        if p is None:
            raise gr.Error("Create a project first (New Project).")
        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 generated by LLM (ZeroGPU).")

    propose_btn.click(
        on_propose,
        inputs=[project, sb_prompt, sb_target_shots, sb_default_fps, sb_default_len],
        outputs=[project, shots_json, sb_status]
    )

    def on_confirm(p):
        if p is None or not p.get("shots"):
            raise gr.Error("No storyboard yet.")
        edges = []
        for i in range(len(p["shots"]) - 1):
            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()
        save_project(p)
        return (
            p,
            gr.update(value=p["shots"]),
            gr.update(value=p["clips"]),
            gr.update(value="Storyboard confirmed. Proceed to Keyframes."),
            gr.update(interactive=True)
        )

    confirm_btn.click(
        on_confirm,
        inputs=[project],
        outputs=[project, kf_table, vd_table, sb_status, to_videos_btn]
    )

    def on_save(p):
        if p is None:
            raise gr.Error("No project in memory.")
        path = save_project(p)
        return gr.update(value=f"Saved to `{path}`")

    save_btn.click(on_save, inputs=[project], outputs=[sb_status])

    def on_load(file_obj):
        p = load_project_file(file_obj)
        return (
            p,
            gr.update(value=f"Loaded project `{p['meta']['name']}` (id: `{p['meta']['id']}`)"),
            gr.update(value=p["shots"]),
            gr.update(value=p["clips"]),
            gr.update(interactive=bool(p.get("shots")))
        )

    load_btn.click(
        on_load,
        inputs=[load_file],
        outputs=[project, sb_status, kf_table, vd_table, to_videos_btn]
    )

if __name__ == "__main__":
    # SSR is fine; you can set share=True if you want a public link automatically
    demo.launch()