Spaces:
Running on Zero
Running on Zero
| import os | |
| import tempfile | |
| from pathlib import Path | |
| # Alguns workers ZeroGPU não expõem NVML completamente ao allocator do PyTorch. | |
| # Desativar o cache evita o assert interno sem alterar o resultado do modelo. | |
| os.environ.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1") | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from diffusers import CogVideoXImageToVideoPipeline | |
| from diffusers.utils import export_to_video | |
| from PIL import Image | |
| MODEL_ID = "THUDM/CogVideoX1.5-5B-I2V" | |
| PIPE = None | |
| def get_pipeline(): | |
| global PIPE | |
| if PIPE is None: | |
| print(f"Loading model: {MODEL_ID}", flush=True) | |
| PIPE = CogVideoXImageToVideoPipeline.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=torch.bfloat16, | |
| ) | |
| # O offload evita mover o pipeline inteiro de uma só vez para a GPU. | |
| PIPE.enable_model_cpu_offload() | |
| print("Model loaded on CUDA", flush=True) | |
| return PIPE | |
| def resize_image(image: Image.Image, profile: str): | |
| image = image.convert("RGB") | |
| if profile == "1080p (upscale final)": | |
| target_w, target_h = 1360, 768 | |
| else: | |
| target_w, target_h = 1024, 576 | |
| target_ratio = target_w / target_h | |
| source_ratio = image.width / image.height | |
| if source_ratio > target_ratio: | |
| crop_h = image.height | |
| crop_w = round(crop_h * target_ratio) | |
| else: | |
| crop_w = image.width | |
| crop_h = round(crop_w / target_ratio) | |
| left = max(0, (image.width - crop_w) // 2) | |
| top = max(0, (image.height - crop_h) // 2) | |
| image = image.crop((left, top, left + crop_w, top + crop_h)) | |
| return image.resize((target_w, target_h), Image.Resampling.LANCZOS), target_w, target_h | |
| def generate_video(image, prompt, profile, seed): | |
| if image is None: | |
| raise gr.Error("Envie uma imagem antes de gerar o vídeo.") | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Escreva o movimento que você quer animar.") | |
| try: | |
| pipeline = get_pipeline() | |
| image, width, height = resize_image(image, profile) | |
| generator = torch.Generator(device="cuda").manual_seed(int(seed or 42)) | |
| full_prompt = ( | |
| "Animate the input image while preserving the exact identity, face, clothing, " | |
| "objects, composition and environment. Focus on natural physically plausible motion. " | |
| + prompt.strip() | |
| + ". Smooth cinematic movement, consistent character, no morphing, no extra limbs, " | |
| "no scene change, no camera cuts." | |
| ) | |
| result = pipeline( | |
| prompt=full_prompt, | |
| image=image, | |
| height=height, | |
| width=width, | |
| num_frames=33, | |
| guidance_scale=6.0, | |
| num_inference_steps=12, | |
| generator=generator, | |
| ) | |
| frames = result.frames[0] | |
| output_path = Path(tempfile.mkdtemp(prefix="motion-ai-")) / "motion.mp4" | |
| export_to_video(frames, str(output_path), fps=8) | |
| return str(output_path) | |
| except Exception as exc: | |
| print(f"Generation error: {type(exc).__name__}: {exc}", flush=True) | |
| raise gr.Error(f"Não foi possível gerar o vídeo: {type(exc).__name__}: {exc}") from exc | |
| with gr.Blocks(title="Motion AI Studio") as demo: | |
| gr.Markdown( | |
| "# Motion AI Studio\n" | |
| "Anime uma imagem com um comando de movimento. A GPU é solicitada somente ao gerar." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(type="pil", label="Imagem de referência") | |
| prompt = gr.Textbox( | |
| label="Descreva o movimento", | |
| placeholder="Ex.: o homem começa a correr pela rua enquanto a câmera o acompanha lateralmente", | |
| lines=4, | |
| ) | |
| profile = gr.Radio( | |
| ["720p", "1080p (upscale final)"], | |
| value="720p", | |
| label="Perfil", | |
| ) | |
| seed = gr.Number(value=42, precision=0, label="Seed") | |
| generate = gr.Button("Gerar animação", variant="primary") | |
| with gr.Column(): | |
| result = gr.Video(label="Resultado", format="mp4") | |
| gr.Markdown( | |
| "**Recomendação:** comece com 720p. O perfil 1080p gera a animação em resolução menor." | |
| ) | |
| generate.click(generate_video, [image, prompt, profile, seed], result) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=8).launch() | |