Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |