# Prediction interface for Cog ⚙️ # https://cog.run/python import os import subprocess import time import torch from diffusers import CogVideoXPipeline from diffusers.utils import export_to_video from cog import BasePredictor, Input, Path MODEL_CACHE = "model_cache" MODEL_URL = f"https://weights.replicate.delivery/default/THUDM/CogVideo/{MODEL_CACHE}.tar" os.environ["HF_DATASETS_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" os.environ["HF_HOME"] = MODEL_CACHE os.environ["TORCH_HOME"] = MODEL_CACHE os.environ["HF_DATASETS_CACHE"] = MODEL_CACHE os.environ["TRANSFORMERS_CACHE"] = MODEL_CACHE os.environ["HUGGINGFACE_HUB_CACHE"] = MODEL_CACHE def download_weights(url, dest): start = time.time() print("downloading url: ", url) print("downloading to: ", dest) subprocess.check_call(["pget", "-x", url, dest], close_fds=False) print("downloading took: ", time.time() - start) class Predictor(BasePredictor): def setup(self) -> None: """Load the model into memory to make running multiple predictions efficient""" if not os.path.exists(MODEL_CACHE): download_weights(MODEL_URL, MODEL_CACHE) # model_id: THUDM/CogVideoX-5b self.pipe = CogVideoXPipeline.from_pretrained( MODEL_CACHE, torch_dtype=torch.bfloat16, ).to("cuda") self.pipe.enable_model_cpu_offload() self.pipe.vae.enable_tiling() def predict( self, prompt: str = Input( description="Input prompt", default="A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance.", ), num_inference_steps: int = Input( description="Number of denoising steps", ge=1, le=500, default=50 ), guidance_scale: float = Input( description="Scale for classifier-free guidance", ge=1, le=20, default=6 ), num_frames: int = Input(description="Number of frames for the output video", default=49), seed: int = Input( description="Random seed. Leave blank to randomize the seed", default=None ), ) -> Path: """Run a single prediction on the model""" if seed is None: seed = int.from_bytes(os.urandom(2), "big") print(f"Using seed: {seed}") video = self.pipe( prompt=prompt, num_videos_per_prompt=1, num_inference_steps=num_inference_steps, num_frames=num_frames, guidance_scale=guidance_scale, generator=torch.Generator(device="cuda").manual_seed(seed), ).frames[0] out_path = "/tmp/out.mp4" export_to_video(video, out_path, fps=8) return Path(out_path)