|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import subprocess |
|
|
import time |
|
|
import torch |
|
|
from diffusers import CogVideoXImageToVideoPipeline |
|
|
from diffusers.utils import export_to_video, load_image |
|
|
from cog import BasePredictor, Input, Path |
|
|
|
|
|
|
|
|
MODEL_CACHE = "model_cache_i2v" |
|
|
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) |
|
|
|
|
|
|
|
|
self.pipe = CogVideoXImageToVideoPipeline.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="Starry sky slowly rotating." |
|
|
), |
|
|
image: Path = Input(description="Input image"), |
|
|
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}") |
|
|
|
|
|
img = load_image(image=str(image)) |
|
|
|
|
|
video = self.pipe( |
|
|
prompt=prompt, |
|
|
image=img, |
|
|
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) |
|
|
|
