app.py

import os
import tempfile
from typing import Optional

import gradio as gr
import torch
from PIL import Image

try:
    import spaces
except Exception:
    spaces = None

from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video

SPACE_ID = os.getenv("SPACE_ID", "").lower()
IS_VIDEO_SPACE = any(k in SPACE_ID for k in ["hunyuanvideo", "wan-2-1"])

IMAGE_MODEL_ID = os.getenv("IMAGE_MODEL_ID", "black-forest-labs/FLUX.1-schnell")
VIDEO_MODEL_ID = os.getenv("VIDEO_MODEL_ID", "damo-vilab/text-to-video-ms-1.7b")

_image_pipe: Optional[DiffusionPipeline] = None
_video_pipe: Optional[DiffusionPipeline] = None


def _device_dtype():
    if torch.cuda.is_available():
        if torch.cuda.get_device_properties(0).major >= 8:
            return "cuda", torch.bfloat16
        return "cuda", torch.float16
    return "cpu", torch.float32


def _load_image_pipe() -> DiffusionPipeline:
    global _image_pipe
    if _image_pipe is None:
        device, dtype = _device_dtype()
        _image_pipe = DiffusionPipeline.from_pretrained(IMAGE_MODEL_ID, torch_dtype=dtype)
        if device == "cuda":
            _image_pipe.enable_model_cpu_offload()
        else:
            _image_pipe.to("cpu")
    return _image_pipe


def _load_video_pipe() -> DiffusionPipeline:
    global _video_pipe
    if _video_pipe is None:
        device, _ = _device_dtype()
        dtype = torch.float16 if device == "cuda" else torch.float32
        _video_pipe = DiffusionPipeline.from_pretrained(VIDEO_MODEL_ID, torch_dtype=dtype)
        if device == "cuda":
            _video_pipe.enable_model_cpu_offload()
        else:
            _video_pipe.to("cpu")
    return _video_pipe


def _gpu_decorator(seconds: int):
    if spaces is not None:
        return spaces.GPU(duration=seconds)

    def _wrap(fn):
        return fn

    return _wrap


@_gpu_decorator(120)
def generate_image(prompt: str, steps: int, guidance_scale: float, seed: int):
    prompt = (prompt or "").strip() or "A cinematic photo of a woman on a beach at sunset"
    pipe = _load_image_pipe()

    gen = torch.Generator(device="cpu").manual_seed(int(seed))
    image: Image.Image = pipe(
        prompt=prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        generator=gen,
        width=1024,
        height=1024,
    ).images[0]
    return image


@_gpu_decorator(240)
def generate_video(prompt: str, steps: int, fps: int, num_frames: int, seed: int):
    prompt = (prompt or "").strip() or "A woman walking on a sunny beach, cinematic shot"
    pipe = _load_video_pipe()

    gen = torch.Generator(device="cpu").manual_seed(int(seed))
    result = pipe(
        prompt,
        num_inference_steps=int(steps),
        num_frames=int(num_frames),
        generator=gen,
    )
    frames = result.frames[0]

    out_path = os.path.join(tempfile.gettempdir(), "generated_video.mp4")
    export_to_video(frames, out_path, fps=int(fps))
    return out_path


def build_ui():
    title = os.getenv("SPACE_TITLE", SPACE_ID.split("/")[-1].replace("-", " ").title() or "AI Generator")

    if IS_VIDEO_SPACE:
        with gr.Blocks(theme=gr.themes.Soft()) as demo:
            gr.Markdown(f"## {title} — AI Video Generation")
            prompt = gr.Textbox(label="Prompt", value="A woman walking on a sunny beach, cinematic shot")
            with gr.Row():
                steps = gr.Slider(8, 40, value=20, step=1, label="Inference steps")
                num_frames = gr.Slider(8, 32, value=16, step=1, label="Frames")
                fps = gr.Slider(4, 16, value=8, step=1, label="FPS")
                seed = gr.Number(value=42, precision=0, label="Seed")
            out = gr.Video(label="Generated video")
            btn = gr.Button("Generate")
            btn.click(generate_video, [prompt, steps, fps, num_frames, seed], [out])
        return demo

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown(f"## {title} — AI Image Generation")
        prompt = gr.Textbox(label="Prompt", value="A cinematic photo of a woman on a beach at sunset")
        with gr.Row():
            steps = gr.Slider(4, 40, value=20, step=1, label="Inference steps")
            guidance = gr.Slider(1.0, 10.0, value=3.5, step=0.1, label="Guidance scale")
            seed = gr.Number(value=42, precision=0, label="Seed")
        out = gr.Image(type="pil", label="Generated image")
        btn = gr.Button("Generate")
        btn.click(generate_image, [prompt, steps, guidance, seed], [out])
    return demo


demo = build_ui()

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