File size: 5,715 Bytes
208b0eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import json
import os
import numpy as np
import torch
from diffusers import (AutoencoderKL, CogVideoXDDIMScheduler, DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
PNDMScheduler)
from transformers import T5EncoderModel, T5Tokenizer
from omegaconf import OmegaConf
from PIL import Image
from cogvideox.models.transformer3d import CogVideoXTransformer3DModel
from cogvideox.models.autoencoder_magvit import AutoencoderKLCogVideoX
from cogvideox.pipeline.pipeline_cogvideox import CogVideoX_Fun_Pipeline
from cogvideox.pipeline.pipeline_cogvideox_inpaint import CogVideoX_Fun_Pipeline_Inpaint
from cogvideox.utils.lora_utils import merge_lora, unmerge_lora
from cogvideox.utils.utils import get_image_to_video_latent, save_videos_grid, ASPECT_RATIO_512, get_closest_ratio, to_pil
from huggingface_hub import HfApi, HfFolder
# Low GPU memory mode
low_gpu_memory_mode = False
# Model loading section
model_id = "/content/model"
transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
model_id, subfolder="transformer", torch_dtype=torch.bfloat16
).to(torch.bfloat16)
vae = AutoencoderKLCogVideoX.from_pretrained(
model_id, subfolder="vae"
).to(torch.bfloat16)
text_encoder = T5EncoderModel.from_pretrained(
model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
)
sampler_dict = {
"Euler": EulerDiscreteScheduler,
"Euler A": EulerAncestralDiscreteScheduler,
"DPM++": DPMSolverMultistepScheduler,
"PNDM": PNDMScheduler,
"DDIM_Cog": CogVideoXDDIMScheduler,
"DDIM_Origin": DDIMScheduler,
}
scheduler = sampler_dict["DPM++"].from_pretrained(model_id, subfolder="scheduler")
# Pipeline setup
if transformer.config.in_channels != vae.config.latent_channels:
pipeline = CogVideoX_Fun_Pipeline_Inpaint.from_pretrained(
model_id, vae=vae, text_encoder=text_encoder,
transformer=transformer, scheduler=scheduler,
torch_dtype=torch.bfloat16
)
else:
pipeline = CogVideoX_Fun_Pipeline.from_pretrained(
model_id, vae=vae, text_encoder=text_encoder,
transformer=transformer, scheduler=scheduler,
torch_dtype=torch.bfloat16
)
if low_gpu_memory_mode:
pipeline.enable_sequential_cpu_offload()
else:
pipeline.enable_model_cpu_offload()
@torch.inference_mode()
def generate(input):
values = input["input"]
prompt = values["prompt"]
negative_prompt = values.get("negative_prompt", "")
guidance_scale = values.get("guidance_scale", 6.0)
seed = values.get("seed", 42)
num_inference_steps = values.get("num_inference_steps", 50)
base_resolution = values.get("base_resolution", 512)
video_length = values.get("video_length", 53)
fps = values.get("fps", 10)
lora_weight = values.get("lora_weight", 1.00)
save_path = "samples"
partial_video_length = values.get("partial_video_length", None)
overlap_video_length = values.get("overlap_video_length", 4)
validation_image_start = values.get("validation_image_start", "asset/1.png")
validation_image_end = values.get("validation_image_end", None)
generator = torch.Generator(device="cuda").manual_seed(seed)
aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
start_img = Image.open(validation_image_start)
original_width, original_height = start_img.size
closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size)
height, width = [int(x / 16) * 16 for x in closest_size]
sample_size = [height, width]
if partial_video_length is not None:
# Handle ultra-long video generation if required
# ... (existing logic for partial video generation)
else:
# Standard video generation
video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
input_video, input_video_mask, clip_image = get_image_to_video_latent(validation_image_start, validation_image_end, video_length=video_length, sample_size=sample_size)
with torch.no_grad():
sample = pipeline(
prompt=prompt,
num_frames=video_length,
negative_prompt=negative_prompt,
height=sample_size[0],
width=sample_size[1],
generator=generator,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
video=input_video,
mask_video=input_video_mask
).videos
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
video_path = os.path.join(save_path, f"{prefix}.mp4")
save_videos_grid(sample, video_path, fps=fps)
# Upload final video to Hugging Face repository
#hf_api = HfApi()
#repo_id = values.get("repo_id", "your-username/your-repo") # Set your HF repo
#hf_api.upload_file(
# path_or_fileobj=video_path,
# path_in_repo=f"{prefix}.mp4",
# repo_id=repo_id,
# repo_type="model" # or "dataset" if using a dataset repo
#)
# Prepare output
#result_url = f"https://huggingface.co/{repo_id}/blob/main/{prefix}.mp4"
result_url = ""
job_id = values.get("job_id", "default-job-id") # For RunPod job tracking
return {"jobId": job_id, "result": result_url, "status": "DONE"}
runpod.serverless.start({"handler": generate})
|