app.py

from __future__ import annotations

from typing import List, Tuple

import gradio as gr

from cinegen import CharacterDesigner, StoryGenerator, VideoDirector
from cinegen.models import Storyboard

try:  # pragma: no cover - spaces is only available inside HF Spaces
    import spaces  # type: ignore
except Exception:  # pragma: no cover - keep local dev working without spaces pkg
    spaces = None  # type: ignore

if spaces:
    @spaces.GPU(duration=60)  # short duration is enough
    def __cinegen_gpu_warmup():
        """Dummy function — never called, only exists to satisfy HF Spaces GPU detection"""
        pass


STYLE_CHOICES = [
    "Cinematic Realism",
    "Neo-Noir Animation",
    "Analog Horror",
    "Retro-Futuristic",
    "Dreamlike Documentary",
]

VIDEO_MODEL_CHOICES = [
    ("Wan 2.2 TI2V (fal-ai)", "Wan-AI/Wan2.2-TI2V-5B"),
    ("LTX Video 0.9.7", "Lightricks/LTX-Video-0.9.7-distilled"),
    ("Hunyuan Video 1.5", "tencent/HunyuanVideo-1.5"),
    ("CogVideoX 5B", "THUDM/CogVideoX-5b"),
]

SCENE_COLUMNS = ["Scene", "Title", "Action", "Visuals", "Characters", "Duration (s)"]
CHARACTER_COLUMNS = ["ID", "Name", "Role", "Traits"]


def gpu_guard(duration: int = 120):
    def decorator(fn):
        if not spaces:
            return fn
        return spaces.GPU(duration=duration)(fn)
    return decorator

def _character_dropdown_update(board: Storyboard | None):
    if not board or not board.characters:
        return gr.update(choices=[], value=None, interactive=False)
    choices = [character.identifier for character in board.characters]
    return gr.update(choices=choices, value=choices[0], interactive=True)


def _gallery_from_board(board: Storyboard) -> List[Tuple[str, str]]:
    gallery: List[Tuple[str, str]] = []
    for character in board.characters:
        if not character.reference_image:
            continue
        caption = f"{character.name} — {character.role}"
        gallery.append((character.reference_image, caption))
    return gallery


def _ensure_storyboard(board: Storyboard | None) -> Storyboard:
    if not board:
        raise gr.Error("Create a storyboard first.")
    return board


def _validate_inputs(idea: str | None, image_path: str | None):
    if not idea and not image_path:
        raise gr.Error("Provide either a story idea or upload a reference image.")


def handle_storyboard(
    idea: str,
    inspiration_image: str | None,
    style: str,
    scene_count: int,
    google_api_key: str,
) -> Tuple[str, List[List[str]], List[List[str]], Storyboard, dict]:
    _validate_inputs(idea, inspiration_image)
    generator = StoryGenerator(api_key=google_api_key or None)
    storyboard = generator.generate(
        idea=idea,
        style=style,
        scene_count=scene_count,
        inspiration_path=inspiration_image,
    )
    summary_md = f"### {storyboard.title}\n{storyboard.synopsis}"
    scene_rows = storyboard.scenes_table()
    character_rows = storyboard.characters_table()
    dropdown_update = _character_dropdown_update(storyboard)
    return (
        summary_md,
        [[row[col] for col in SCENE_COLUMNS] for row in scene_rows],
        [[row[col] for col in CHARACTER_COLUMNS] for row in character_rows],
        storyboard,
        dropdown_update,
    )


def handle_character_design(
    storyboard: Storyboard | None,
    google_api_key: str,
):
    board = _ensure_storyboard(storyboard)
    designer = CharacterDesigner(api_key=google_api_key or None)
    _, updated_board = designer.design(board)
    gallery = _gallery_from_board(updated_board)
    if not gallery:
        raise gr.Error("Failed to design characters.")
    return gallery, updated_board


def handle_character_regen(
    storyboard: Storyboard | None,
    character_id: str | None,
    google_api_key: str,
):
    board = _ensure_storyboard(storyboard)
    if not character_id:
        raise gr.Error("Select a character ID to regenerate.")
    designer = CharacterDesigner(api_key=google_api_key or None)
    try:
        _, updated_board = designer.redesign_character(board, character_id)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc
    gallery = _gallery_from_board(updated_board)
    if not gallery:
        raise gr.Error("Failed to refresh character art.")
    return gallery, updated_board


@gpu_guard(duration=300)
def handle_video_render(
    storyboard: Storyboard | None,
    hf_token: str,
    model_choice: str,
):
    board = _ensure_storyboard(storyboard)
    prioritized_models = [model_choice] + [
        model for _, model in VIDEO_MODEL_CHOICES if model != model_choice
    ]
    director = VideoDirector(token=hf_token or None, models=prioritized_models)
    final_cut, logs = director.render(board)
    log_md = "\n".join(f"- {line}" for line in logs)
    return final_cut, log_md


css = """
#cinegen-app {
    max-width: 1080px;
    margin: 0 auto;
}
"""


with gr.Blocks(fill_height=True, elem_id="cinegen-app") as demo:
    gr.Markdown(
        "## 🎬 CineGen AI Director\n"
        "Drop an idea or inspiration image and let CineGen produce a storyboard, character boards, "
        "and a compiled short film using Hugging Face video models."
    )

    story_state = gr.State()

    with gr.Row():
        idea_box = gr.Textbox(
            label="Movie Idea",
            placeholder="E.g. A time loop love story set in a neon bazaar.",
            lines=3,
        )
        inspiration = gr.Image(label="Reference Image (optional)", type="filepath")

    with gr.Row():
        style_dropdown = gr.Dropdown(
            label="Visual Style",
            choices=STYLE_CHOICES,
            value=STYLE_CHOICES[0],
        )
        scene_slider = gr.Slider(
            label="Scene Count",
            minimum=3,
            maximum=8,
            value=4,
            step=1,
        )
        video_model_dropdown = gr.Dropdown(
            label="Preferred Video Model",
            choices=[choice for choice, _ in VIDEO_MODEL_CHOICES],
            value=VIDEO_MODEL_CHOICES[0][0],
        )

    with gr.Accordion("API Keys", open=True):
        gr.Markdown(
            "Provide your own API credentials for live Gemini and Hugging Face calls. "
            "Keys stay within your browser session and are not stored on the server."
        )
        google_key_input = gr.Textbox(
            label="Google API Key (Gemini)",
            type="password",
            placeholder="Required for live Gemini calls. Leave blank to use offline stubs.",
        )
        hf_token_input = gr.Textbox(
            label="Hugging Face Token",
            type="password",
            placeholder="Needed for Wan/LTX/Hunyuan video generation.",
        )

    storyboard_btn = gr.Button("Create Storyboard", variant="primary")
    summary_md = gr.Markdown("Storyboard output will appear here.")
    scenes_df = gr.Dataframe(headers=SCENE_COLUMNS, wrap=True)
    characters_df = gr.Dataframe(headers=CHARACTER_COLUMNS, wrap=True)

    with gr.Row():
        design_btn = gr.Button("Design Characters", variant="secondary")
        render_btn = gr.Button("Render Short Film", variant="primary")

    with gr.Row():
        character_select = gr.Dropdown(
            label="Character Slot",
            choices=[],
            interactive=False,
            info="Select an ID from the storyboard table to regenerate its portrait.",
        )
        regen_btn = gr.Button("Regenerate Selected Character", variant="secondary")

    gallery = gr.Gallery(label="Character References", columns=4, height=320)
    render_logs = gr.Markdown(label="Render Log")
    final_video = gr.Video(label="CineGen Short Film", interactive=False)

    storyboard_btn.click(
        fn=handle_storyboard,
        inputs=[idea_box, inspiration, style_dropdown, scene_slider, google_key_input],
        outputs=[summary_md, scenes_df, characters_df, story_state, character_select],
    )

    design_btn.click(
        fn=handle_character_design,
        inputs=[story_state, google_key_input],
        outputs=[gallery, story_state],
    )

    regen_btn.click(
        fn=handle_character_regen,
        inputs=[story_state, character_select, google_key_input],
        outputs=[gallery, story_state],
    )

    def _model_value(label: str) -> str:
        lookup = dict(VIDEO_MODEL_CHOICES)
        return lookup.get(label, VIDEO_MODEL_CHOICES[0][1])

    def render_wrapper(board, token, label):
        return handle_video_render(board, token, _model_value(label))

    render_btn.click(
        fn=render_wrapper,
        inputs=[story_state, hf_token_input, video_model_dropdown],
        outputs=[final_video, render_logs],
        queue=True,
        concurrency_limit=1,
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), css=css)