app.py

"""Hugging Face Space entry point — venture-studio deployed as a ZeroGPU Space.

Push this whole repo to a HF Space (sdk: gradio). The Space's Python build will
install `requirements.txt` (torch + diffusers + transformers + gradio + spaces)
and call:
  - `app.py:infer`         — single image → motion (AnimateDiff MotionLoRA)
  - `app.py:infer_txt2img` — prompt → 512×512 sprite (SD 1.5)
  - `app.py:infer_ltx_i2v` — reference portrait → portrait video (LTX-Video I2V)

`@spaces.GPU` allocates an A10G only for the duration of the call (ZeroGPU
model), so the Space is free for the maintainer and shared fairly across users.

Locally this file is not imported — `studio` CLI uses `studio.backends.hf_space`
to call this Space remotely via gradio_client.
"""
from __future__ import annotations
import io
import tempfile

import gradio as gr
import numpy as np
import spaces
from PIL import Image as PILImage

from pixel_cursor import open_cursor, FrameStack
from pixel_cursor.artifact import _new_image
from studio.backends.animatediff import AnimateDiffAdapter, MOTION_LORA_MAP


_adapter: AnimateDiffAdapter | None = None
_sd_pipe = None
_ltx_pipe = None
_wan_pixel_pipe = None
_asr_pipe = None
_diar_pipe = None


def _get_adapter() -> AnimateDiffAdapter:
    global _adapter
    if _adapter is None:
        _adapter = AnimateDiffAdapter()
        _adapter.register()
    return _adapter


PIXEL_ART_LORA_REPO = "artificialguybr/pixelartredmond-1-5v-pixel-art-loras-for-sd-1-5"
PIXEL_ART_LORA_WEIGHT_FILE = "PixelArtRedmond15V-PixelArt-PIXARFK.safetensors"
PIXEL_ART_LORA_ADAPTER = "pixart"


def _get_sd_pipe():
    global _sd_pipe
    if _sd_pipe is not None:
        return _sd_pipe
    import torch
    from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler
    pipe = AutoPipelineForText2Image.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16,
        safety_checker=None,
        requires_safety_checker=False,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")
    pipe.set_progress_bar_config(disable=True)
    try:
        pipe.load_lora_weights(
            PIXEL_ART_LORA_REPO,
            weight_name=PIXEL_ART_LORA_WEIGHT_FILE,
            adapter_name=PIXEL_ART_LORA_ADAPTER,
        )
        pipe.set_adapters([PIXEL_ART_LORA_ADAPTER], adapter_weights=[0.9])
        print(f"loaded pixel-art LoRA: {PIXEL_ART_LORA_REPO}", flush=True)
    except Exception as e:
        print(f"WARN: could not load LoRA {PIXEL_ART_LORA_REPO}: {e}", flush=True)
    _sd_pipe = pipe
    return pipe


def _get_ltx_pipe():
    """Lazily load LTX-Video I2V pipeline (bfloat16, CUDA).

    Model: Lightricks/LTX-Video (~8 GB, cached in persistent storage after first call).
    Supports portrait aspect ratios (height > width) with both dims divisible by 32.
    Frame counts must be of the form 8k+1 (9, 17, 25, 49, 97, 121 ...).
    """
    global _ltx_pipe
    if _ltx_pipe is not None:
        return _ltx_pipe
    import torch
    from diffusers import LTXImageToVideoPipeline
    pipe = LTXImageToVideoPipeline.from_pretrained(
        "Lightricks/LTX-Video",
        torch_dtype=torch.bfloat16,
    )
    pipe = pipe.to("cuda")
    pipe.set_progress_bar_config(disable=True)
    print("LTX-Video I2V pipeline loaded", flush=True)
    _ltx_pipe = pipe
    return pipe


def _get_wan_pixel_pipe():
    """Load Wan 2.2 I2V with the pixel-sprite animation LoRA."""
    global _wan_pixel_pipe
    if _wan_pixel_pipe is not None:
        return _wan_pixel_pipe
    import torch
    from diffusers import DiffusionPipeline

    pipe = DiffusionPipeline.from_pretrained(
        "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    )
    pipe.load_lora_weights(
        "styly-agents/Wan2-2-pixel-animate",
        weight_name="wan2.2_animate_adapter_model.safetensors",
        adapter_name="pixel_animate",
    )
    pipe.set_adapters(["pixel_animate"], adapter_weights=[1.0])
    pipe.set_progress_bar_config(disable=True)
    print("Wan 2.2 pixel animation adapter loaded", flush=True)
    _wan_pixel_pipe = pipe
    return pipe


def _get_asr_pipe():
    """Lazily load a GPU Whisper ASR pipeline (transformers, already a dependency)."""
    global _asr_pipe
    if _asr_pipe is not None:
        return _asr_pipe
    import torch
    from transformers import pipeline as hf_pipeline

    _asr_pipe = hf_pipeline(
        "automatic-speech-recognition",
        model="openai/whisper-small",
        torch_dtype=torch.float16,
        device="cuda",
        return_timestamps=True,
    )
    print("Whisper (small) ASR pipeline loaded on cuda", flush=True)
    return _asr_pipe


def _get_diar_pipe():
    """Lazily load the pyannote speaker-diarization pipeline.

    Requires the `HF_TOKEN` Space secret (set via the Spaces UI/API, never
    committed to this repo) — the token must have accepted the gated terms
    for pyannote/speaker-diarization-3.1 and pyannote/segmentation-3.0, and
    have "read access to public gated repos" enabled in its token settings.
    """
    global _diar_pipe
    if _diar_pipe is not None:
        return _diar_pipe
    import os
    from pyannote.audio import Pipeline as DiarPipeline

    token = os.environ.get("HF_TOKEN")
    _diar_pipe = DiarPipeline.from_pretrained(
        "pyannote/speaker-diarization-3.1", token=token
    )
    _diar_pipe.to_device = getattr(_diar_pipe, "to_device", None)
    try:
        import torch
        _diar_pipe.to(torch.device("cuda"))
    except Exception as e:
        print(f"WARN: could not move diarization pipeline to cuda: {e}", flush=True)
    print("pyannote speaker-diarization-3.1 pipeline loaded", flush=True)
    return _diar_pipe


@spaces.GPU(duration=300)
def infer_audio(audio_path: str) -> tuple[str, str]:
    """Transcribe + diarize an uploaded audio/video file.

    Returns (transcript_text, diarization_json_text). Transcript segments are
    timestamped; diarization turns are (start, end, speaker_label) — labels
    are anonymous (SPEAKER_00, SPEAKER_01, ...), not real identities.
    """
    import json

    asr = _get_asr_pipe()
    asr_result = asr(
        audio_path,
        chunk_length_s=30,
        batch_size=8,
        generate_kwargs={"language": "en"},
    )

    lines = []
    for chunk in asr_result.get("chunks", []):
        start, end = chunk.get("timestamp", (None, None))
        text = chunk.get("text", "").strip()
        if start is None:
            lines.append(text)
        else:
            lines.append(f"[{start:07.2f} -> {end:07.2f}] {text}")
    transcript_text = "\n".join(lines) if lines else asr_result.get("text", "")

    try:
        diar = _get_diar_pipe()
        diarization = diar(audio_path)
        turns = [
            {"start": round(turn.start, 2), "end": round(turn.end, 2), "speaker": speaker}
            for turn, _, speaker in diarization.itertracks(yield_label=True)
        ]
        diarization_json = json.dumps(turns, indent=2)
    except Exception as e:
        # pyannote 3.4.0 (required for the gated 3.1 pipeline) needs an older
        # torchaudio than this Space's torch stack ships (needed by Wan2.2 /
        # LTX-Video / spaces itself) — diarization is best-effort here rather
        # than something allowed to take down the whole call.
        print(f"WARN: diarization unavailable: {e}", flush=True)
        diarization_json = json.dumps(
            {"error": "diarization unavailable on this Space (torch version conflict with pyannote)"}
        )

    return transcript_text, diarization_json


@spaces.GPU(duration=90)
def infer(
    image: np.ndarray,
    preset: str,
    num_frames: int,
    num_inference_steps: int,
    guidance_scale: float,
    prompt: str,
    negative_prompt: str,
    seed: int,
) -> str:
    """Generate motion on a single image. Returns path to an mp4 file."""
    import tempfile
    import imageio.v3 as iio

    adapter = _get_adapter()
    img = _new_image(image.astype(np.uint8))
    cur = open_cursor(img).bind_motion_exemplar([preset], backend="animatediff")

    motion_spec = {
        "num_frames": int(num_frames),
        "num_inference_steps": int(num_inference_steps),
        "guidance_scale": float(guidance_scale),
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "seed": int(seed),
    }
    stack: FrameStack = cur.write_motion(motion_spec, backend="animatediff")

    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    iio.imwrite(out.name, stack.frames, fps=stack.fps)
    return out.name


@spaces.GPU(duration=45)
def infer_txt2img(
    prompt: str,
    negative_prompt: str,
    num_inference_steps: int,
    guidance_scale: float,
    height: int,
    width: int,
    seed: int,
    lora_weight: float,
) -> str:
    """Generate a single sprite from a text prompt. Returns path to a PNG."""
    import torch

    pipe = _get_sd_pipe()
    try:
        pipe.set_adapters([PIXEL_ART_LORA_ADAPTER], adapter_weights=[float(lora_weight)])
    except Exception as e:
        print(f"WARN: set_adapters failed: {e}", flush=True)
    g = torch.Generator(device="cuda").manual_seed(int(seed))
    out = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        height=int(height),
        width=int(width),
        generator=g,
    )
    image = out.images[0]
    f = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    image.save(f.name)
    return f.name


@spaces.GPU(duration=150)
def infer_ltx_i2v(
    image: np.ndarray,
    prompt: str,
    negative_prompt: str,
    height: int,
    width: int,
    num_frames: int,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
) -> str:
    """LTX-Video image-to-video: reference portrait → portrait video clip.

    Constraints enforced here (not in UI) so programmatic callers are safe:
      - height and width are rounded up to nearest multiple of 32
      - num_frames is rounded up to nearest 8k+1 value

    Returns path to an mp4 file at 24fps.
    """
    import torch
    import imageio.v3 as iio

    # Enforce divisibility constraints
    h = int(height)
    w = int(width)
    h = ((h + 31) // 32) * 32
    w = ((w + 31) // 32) * 32
    nf = int(num_frames)
    if (nf - 1) % 8 != 0:
        nf = ((nf // 8) * 8) + 1

    pipe = _get_ltx_pipe()
    pil_img = PILImage.fromarray(image.astype(np.uint8))
    gen = torch.Generator(device="cuda").manual_seed(int(seed))

    result = pipe(
        image=pil_img,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=h,
        width=w,
        num_frames=nf,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        generator=gen,
    )
    frames_pil = result.frames[0]
    frames_np = np.stack([np.array(f) for f in frames_pil], axis=0)

    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    iio.imwrite(out.name, frames_np, fps=24)
    return out.name


@spaces.GPU(duration=180)
def infer_wan_pixel(
    image: np.ndarray,
    prompt: str,
    negative_prompt: str,
    height: int,
    width: int,
    num_frames: int,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
) -> str:
    """Generate identity-focused pixel sprite motion with Wan 2.2 + LoRA."""
    import torch
    from diffusers.utils import export_to_video

    h = max(256, min(480, int(height)))
    w = max(256, min(832, int(width)))
    h = (h // 16) * 16
    w = (w // 16) * 16
    nf = max(8, min(32, int(num_frames)))
    pipe = _get_wan_pixel_pipe()
    pil_img = PILImage.fromarray(image.astype(np.uint8))
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    result = pipe(
        image=pil_img,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=h,
        width=w,
        num_frames=nf,
        num_inference_steps=max(4, min(20, int(num_inference_steps))),
        guidance_scale=float(guidance_scale),
        generator=generator,
    )
    frames = result.frames[0]
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    export_to_video(frames, out.name, fps=16)
    return out.name


with gr.Blocks(title="Venture-Studio") as demo:
    gr.Markdown(
        "# Venture-Studio · Pixel-Cursor Animation\n"
        "Single image → 24fps animation, text prompt → sprite, or portrait → locked video.\n"
        "Running on Hugging Face ZeroGPU."
    )
    with gr.Tabs():
        with gr.Tab("Motion"):
            with gr.Row():
                with gr.Column():
                    image_in = gr.Image(label="Source image", type="numpy", height=384)
                    preset = gr.Dropdown(
                        choices=sorted(MOTION_LORA_MAP), value="zoom_in",
                        label="MotionLoRA preset",
                    )
                    with gr.Accordion("Advanced", open=False):
                        num_frames = gr.Slider(8, 24, value=16, step=2, label="num_frames")
                        steps = gr.Slider(10, 50, value=25, step=1, label="num_inference_steps")
                        guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="guidance_scale")
                        prompt = gr.Textbox(value="high quality, detailed", label="prompt")
                        neg = gr.Textbox(value="bad quality, blurry", label="negative_prompt")
                        seed = gr.Number(value=42, precision=0, label="seed")
                    run = gr.Button("Generate", variant="primary")
                with gr.Column():
                    video_out = gr.Video(label="Output", autoplay=True, loop=True)

            run.click(
                infer,
                inputs=[image_in, preset, num_frames, steps, guidance, prompt, neg, seed],
                outputs=video_out,
                api_name="infer",
            )

        with gr.Tab("Sprite gen (txt2img)"):
            with gr.Row():
                with gr.Column():
                    t2i_prompt = gr.Textbox(
                        value=(
                            "pixel art, PixArFK, fantasy goblin warrior, green skin, "
                            "leather armor, empty hands, unarmed, standing pose, "
                            "full body, centered, white background, retro game sprite"
                        ),
                        lines=3, label="prompt (include 'pixel art, PixArFK' for LoRA)",
                    )
                    t2i_neg = gr.Textbox(
                        value=(
                            "sword, weapon, dagger, axe, staff, blurry, soft, "
                            "photorealistic, 3d render, extra limbs, distorted, "
                            "multiple characters"
                        ),
                        lines=2, label="negative_prompt",
                    )
                    with gr.Accordion("Advanced", open=False):
                        t2i_steps = gr.Slider(10, 50, value=25, step=1, label="num_inference_steps")
                        t2i_guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="guidance_scale")
                        t2i_height = gr.Slider(256, 768, value=512, step=64, label="height")
                        t2i_width = gr.Slider(256, 768, value=512, step=64, label="width")
                        t2i_seed = gr.Number(value=0, precision=0, label="seed (0 = random)")
                        t2i_lora = gr.Slider(0.0, 1.5, value=0.9, step=0.05, label="LoRA weight (PixelArtRedmond)")
                    t2i_run = gr.Button("Generate sprite", variant="primary")
                with gr.Column():
                    t2i_out = gr.Image(label="Generated sprite", height=512)

            t2i_run.click(
                infer_txt2img,
                inputs=[t2i_prompt, t2i_neg, t2i_steps, t2i_guidance,
                        t2i_height, t2i_width, t2i_seed, t2i_lora],
                outputs=t2i_out,
                api_name="infer_txt2img",
            )

        with gr.Tab("Hologram (LTX I2V)"):
            gr.Markdown(
                "### LTX-Video Image-to-Video\n"
                "Reference portrait image → locked-head portrait video clip. "
                "Height must exceed width (portrait). Both dims rounded to nearest 32. "
                "Frames rounded to nearest 8k+1 (9 17 25 49 97 121...)."
            )
            with gr.Row():
                with gr.Column():
                    ltx_image = gr.Image(label="Reference portrait", type="numpy", height=384)
                    ltx_prompt = gr.Textbox(
                        value=(
                            "East-Asian man, 30s, dark navy suit, subtle lapel pin, "
                            "neutral-formal expression, studio lighting, soft rim light, "
                            "portrait frame, subject upper two-thirds of frame. "
                            "Speaking naturally, lips 60-70% open, visible lip movement. "
                            "Head absolutely still, zero lateral or vertical drift. "
                            "Single continuous shot, no cuts. Photorealistic, cinematic."
                        ),
                        lines=4, label="prompt",
                    )
                    ltx_neg = gr.Textbox(
                        value=(
                            "head movement, swaying, bobbing, nodding, camera shake, "
                            "zoom, pan, closed mouth, jump cut, cartoon, deformed, blurry"
                        ),
                        lines=2, label="negative_prompt",
                    )
                    with gr.Accordion("Advanced", open=False):
                        ltx_height = gr.Slider(256, 768, value=576, step=32,
                                               label="height (portrait: height > width, div-32)")
                        ltx_width = gr.Slider(256, 768, value=320, step=32,
                                              label="width (div-32)")
                        ltx_frames = gr.Slider(9, 121, value=121, step=8,
                                               label="num_frames (8k+1: 9 17 25 49 97 121)")
                        ltx_steps = gr.Slider(10, 50, value=25, step=1, label="num_inference_steps")
                        ltx_guidance = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="guidance_scale")
                        ltx_seed = gr.Number(value=42, precision=0, label="seed")
                    ltx_run = gr.Button("Generate", variant="primary")
                with gr.Column():
                    ltx_out = gr.Video(label="Output", autoplay=True, loop=True)

            ltx_run.click(
                infer_ltx_i2v,
                inputs=[ltx_image, ltx_prompt, ltx_neg,
                        ltx_height, ltx_width, ltx_frames,
                        ltx_steps, ltx_guidance, ltx_seed],
                outputs=ltx_out,
                api_name="infer_ltx_i2v",
            )

        with gr.Tab("Pixel animate (Wan 2.2)"):
            gr.Markdown(
                "### Wan 2.2 Pixel Animate\n"
                "Image-to-video sprite animation using the pixel-specific LoRA. "
                "Designed for idle, walk, attack, and VFX motion."
            )
            with gr.Row():
                with gr.Column():
                    wan_image = gr.Image(label="Source sprite", type="numpy", height=384)
                    wan_prompt = gr.Textbox(
                        value=(
                            "pixel art sprite animation, preserve the exact character identity, "
                            "silhouette, palette, and framing; a readable idle animation with "
                            "subtle breathing and cloth motion, crisp edges, stable temporal motion"
                        ),
                        lines=4, label="prompt",
                    )
                    wan_neg = gr.Textbox(
                        value="photorealistic, blurry, morphing, extra limbs, camera movement, text, watermark",
                        lines=2, label="negative_prompt",
                    )
                    with gr.Accordion("Advanced", open=False):
                        wan_height = gr.Slider(256, 480, value=368, step=16, label="height")
                        wan_width = gr.Slider(256, 832, value=600, step=16, label="width")
                        wan_frames = gr.Slider(8, 32, value=16, step=8, label="num_frames")
                        wan_steps = gr.Slider(4, 20, value=8, step=1, label="num_inference_steps")
                        wan_guidance = gr.Slider(1.0, 6.0, value=1.1, step=0.1, label="guidance_scale")
                        wan_seed = gr.Number(value=42, precision=0, label="seed")
                    wan_run = gr.Button("Animate sprite", variant="primary")
                with gr.Column():
                    wan_out = gr.Video(label="Output", autoplay=True, loop=True)

            wan_run.click(
                infer_wan_pixel,
                inputs=[wan_image, wan_prompt, wan_neg,
                        wan_height, wan_width, wan_frames,
                        wan_steps, wan_guidance, wan_seed],
                outputs=wan_out,
                api_name="infer_wan_pixel",
            )

        with gr.Tab("Audio (Whisper + Diarization)"):
            gr.Markdown(
                "### Whisper transcription + pyannote speaker diarization\n"
                "Upload an audio or video file. Returns a timestamped transcript "
                "and anonymous speaker turns (SPEAKER_00, SPEAKER_01, ...) — this "
                "does NOT identify real names, only separates who-spoke-when."
            )
            with gr.Row():
                with gr.Column():
                    audio_in = gr.Audio(label="Source audio/video", type="filepath")
                    audio_run = gr.Button("Transcribe + diarize", variant="primary")
                with gr.Column():
                    transcript_out = gr.Textbox(
                        label="Timestamped transcript", lines=20
                    )
                    diarization_out = gr.Textbox(
                        label="Speaker turns (JSON)", lines=20
                    )

            audio_run.click(
                infer_audio,
                inputs=[audio_in],
                outputs=[transcript_out, diarization_out],
                api_name="infer_audio",
            )


if __name__ == "__main__":
    demo.launch()
# redeploy trigger 1785341723