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


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


@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",
            )


if __name__ == "__main__":
    demo.launch()