import os import sys import time import shutil import subprocess import spaces from pathlib import Path import gradio as gr from huggingface_hub import snapshot_download APP_DIR = Path(__file__).resolve().parent WORK_DIR = Path(os.getenv("POCKETSORA_WORKDIR", "/data/pocketsora")) REPO_DIR = WORK_DIR / "Open-Sora" CKPT_DIR = REPO_DIR / "ckpts" OUTPUT_DIR = WORK_DIR / "outputs" MODEL_REPO = os.getenv("OPEN_SORA_MODEL_REPO", "hpcai-tech/Open-Sora-v2") for directory in (WORK_DIR, OUTPUT_DIR): directory.mkdir(parents=True, exist_ok=True) def run_command(command, cwd=None, timeout=None): result = subprocess.run( command, cwd=str(cwd) if cwd else None, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=timeout, env={**os.environ, "PYTHONUNBUFFERED": "1"}, ) if result.returncode != 0: tail = result.stdout[-12000:] if result.stdout else "(sem logs)" raise RuntimeError(f"Comando falhou ({result.returncode}):\n{tail}") return result.stdout or "OK" def prepare_engine(progress=gr.Progress()): """Clone Open-Sora and download the official Open-Sora 2.0 checkpoint.""" if not shutil.which("git"): raise RuntimeError("Git não está instalado no ambiente do Space.") if not (REPO_DIR / "scripts" / "diffusion" / "inference.py").exists(): progress(0.05, desc="Baixando código oficial Open-Sora…") WORK_DIR.mkdir(parents=True, exist_ok=True) run_command([ "git", "clone", "--depth", "1", "https://github.com/hpcaitech/Open-Sora.git", str(REPO_DIR) ]) progress(0.20, desc="Instalando dependências do Open-Sora…") # Official project instructions: install the repository package. marker = REPO_DIR / ".pocketsora_installed" if not marker.exists(): log = run_command([sys.executable, "-m", "pip", "install", "-v", "."], cwd=REPO_DIR, timeout=7200) marker.write_text("installed") else: log = "Dependências do Open-Sora já instaladas nesta sessão." progress(0.45, desc="Baixando checkpoint oficial (download grande)…") CKPT_DIR.mkdir(parents=True, exist_ok=True) snapshot_download( repo_id=MODEL_REPO, local_dir=str(CKPT_DIR), local_dir_use_symlinks=False, resume_download=True, ) expected = [ CKPT_DIR / "Open_Sora_v2.safetensors", CKPT_DIR / "hunyuan_vae.safetensors", ] missing = [str(p) for p in expected if not p.exists()] if missing: raise RuntimeError( "O repositório do modelo foi baixado, mas faltam arquivos esperados:\n" + "\n".join(missing) + "\nConfira a estrutura atual do checkpoint hpcai-tech/Open-Sora-v2." ) progress(0.70, desc="Preparando ambiente de inferência…") return log @spaces.GPU() def generate(prompt, mode, image, aspect_ratio, frames, steps, motion_score, seed, progress=gr.Progress()): prompt = (prompt or "").strip() if not prompt: raise gr.Error("Digite um prompt.") if mode == "Image-to-Video" and image is None: raise gr.Error("Envie uma imagem para Image-to-Video.") try: prepare_engine(progress) job_dir = OUTPUT_DIR / f"job_{int(time.time())}" job_dir.mkdir(parents=True, exist_ok=True) if mode == "Image-to-Video": ref_path = job_dir / "reference.png" image.save(ref_path) config = "configs/diffusion/inference/256px.py" args = [ "--cond_type", "i2v_head", "--ref", str(ref_path), "--prompt", prompt, ] else: # Direct text-to-video uses the official 256px inference config. config = "configs/diffusion/inference/256px.py" args = ["--prompt", prompt] command = [ "torchrun", "--nproc_per_node", "1", "--standalone", "scripts/diffusion/inference.py", config, "--save-dir", str(job_dir), "--aspect_ratio", aspect_ratio, "--num_frames", str(int(frames)), "--num_steps", str(int(steps)), "--motion-score", str(int(motion_score)), "--seed", str(int(seed)), *args, ] progress(0.78, desc="Gerando vídeo com Open-Sora 2.0…") env = { **os.environ, "PYTHONPATH": str(REPO_DIR) + os.pathsep + os.environ.get("PYTHONPATH", ""), "PYTHONUNBUFFERED": "1", } result = subprocess.run( command, cwd=str(REPO_DIR), env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=6 * 60 * 60, ) log = result.stdout or "" if result.returncode != 0: raise RuntimeError(log[-14000:] or f"Open-Sora terminou com código {result.returncode}.") videos = sorted(job_dir.rglob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True) if not videos: raise RuntimeError("A inferência terminou, mas nenhum arquivo MP4 foi encontrado.\n" + log[-6000:]) progress(1.0, desc="Vídeo pronto!") return str(videos[0]), "Concluído.\n\n" + log[-5000:] except subprocess.TimeoutExpired: raise gr.Error("A geração excedeu o limite de 6 horas.") except Exception as exc: raise gr.Error(str(exc)) def system_status(): try: import torch cuda = torch.cuda.is_available() gpu = torch.cuda.get_device_name(0) if cuda else "Nenhuma GPU CUDA detectada" memory = "" if cuda: memory = f"\nVRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB" return f"CUDA: {cuda}\nGPU: {gpu}{memory}\nDiretório de trabalho: {WORK_DIR}" except Exception as exc: return f"Não foi possível consultar PyTorch/CUDA: {exc}" with gr.Blocks(title="PocketSora — Open-Sora 2.0", theme=gr.themes.Soft()) as demo: gr.Markdown( "# PocketSora\n" "### Open-Sora 2.0 · Text-to-Video e Image-to-Video\n" "Interface Gradio para executar o modelo oficial Open-Sora 2.0." ) with gr.Row(): with gr.Column(): mode = gr.Radio( ["Text-to-Video", "Image-to-Video"], value="Text-to-Video", label="Modo de geração", ) prompt = gr.Textbox( label="Prompt", placeholder="A white cartoon cat running through a neon city at night…", lines=4, ) image = gr.Image(label="Imagem de referência (I2V)", type="pil") aspect = gr.Dropdown( ["16:9", "9:16", "1:1", "2.39:1"], value="16:9", label="Proporção", ) frames = gr.Slider(9, 129, value=49, step=4, label="Frames (4k+1)") steps = gr.Slider(10, 50, value=30, step=1, label="Passos") motion = gr.Slider(1, 7, value=4, step=1, label="Motion score") seed = gr.Number(value=42, precision=0, label="Seed") submit = gr.Button("Gerar vídeo", variant="primary") with gr.Column(): output = gr.Video(label="Vídeo gerado", format="mp4") logs = gr.Textbox(label="Logs", lines=14, interactive=False) submit.click( generate, inputs=[prompt, mode, image, aspect, frames, steps, motion, seed], outputs=[output, logs], api_name="generate", ) gr.Markdown("### Diagnóstico do ambiente") status = gr.Textbox(value=system_status(), label="GPU / CUDA", lines=4, interactive=False) refresh = gr.Button("Atualizar diagnóstico") refresh.click(system_status, outputs=status, api_name="status") gr.Markdown( "**Aviso:** Open-Sora 2.0 é um modelo 11B pesado. O benchmark oficial reporta " "52,5 GB de VRAM de pico para 256px em uma H100/H800 com offload. " "Uma GPU pequena, como T4 de 16 GB, não é suficiente para essa configuração." ) if __name__ == "__main__": demo.queue(max_size=2).launch(server_name="0.0.0.0", server_port=7860)