zparadox/stable-video-diffusion
0
1import math2import os3from glob import glob4from pathlib import Path5from typing import Optional6 7import cv28import numpy as np9import torch10from einops import rearrange, repeat11from fire import Fire12from omegaconf import OmegaConf13from PIL import Image14from torchvision.transforms import ToTensor15 16from scripts.util.detection.nsfw_and_watermark_dectection import \17 DeepFloydDataFiltering18from sgm.inference.helpers import embed_watermark19from sgm.util import default, instantiate_from_config20 21import gradio as gr22import uuid23import random24from huggingface_hub import hf_hub_download25 26hf_hub_download(repo_id="stabilityai/stable-video-diffusion-img2vid-xt", filename="svd_xt.safetensors", local_dir="checkpoints") 27 28version = "svd_xt"29device = "cuda"30max_64_bit_int = 2**63 - 131 32def load_model(33 config: str,34 device: str,35 num_frames: int,36 num_steps: int,37):38 config = OmegaConf.load(config)39 if device == "cuda":40 config.model.params.conditioner_config.params.emb_models[41 042 ].params.open_clip_embedding_config.params.init_device = device43 44 config.model.params.sampler_config.params.num_steps = num_steps45 config.model.params.sampler_config.params.guider_config.params.num_frames = (46 num_frames47 )48 if device == "cuda":49 with torch.device(device):50 model = instantiate_from_config(config.model).to(device).eval()51 else:52 model = instantiate_from_config(config.model).to(device).eval()53 54 filter = DeepFloydDataFiltering(verbose=False, device=device)55 return model, filter56 57if version == "svd_xt":58 num_frames = 2559 num_steps = 3060 model_config = "scripts/sampling/configs/svd_xt.yaml"61else:62 raise ValueError(f"Version {version} does not exist.")63 64model, filter = load_model(65 model_config,66 device,67 num_frames,68 num_steps,69)70 71def sample(72 input_path: str = "assets/test_image.png", # Can either be image file or folder with image files73 seed: Optional[int] = None,74 randomize_seed: bool = True,75 motion_bucket_id: int = 127,76 fps_id: int = 6,77 version: str = "svd_xt",78 cond_aug: float = 0.02,79 decoding_t: int = 7, # Number of frames decoded at a time! This eats most VRAM. Reduce if necessary.80 device: str = "cuda",81 output_folder: str = "outputs",82 progress=gr.Progress(track_tqdm=True)83):84 """85 Simple script to generate a single sample conditioned on an image `input_path` or multiple images, one for each86 image file in folder `input_path`. If you run out of VRAM, try decreasing `decoding_t`.87 """88 if(randomize_seed):89 seed = random.randint(0, max_64_bit_int)90 91 torch.manual_seed(seed)92 93 path = Path(input_path)94 all_img_paths = []95 if path.is_file():96 if any([input_path.endswith(x) for x in ["jpg", "jpeg", "png"]]):97 all_img_paths = [input_path]98 else:99 raise ValueError("Path is not valid image file.")100 elif path.is_dir():101 all_img_paths = sorted(102 [103 f104 for f in path.iterdir()105 if f.is_file() and f.suffix.lower() in [".jpg", ".jpeg", ".png"]106 ]107 )108 if len(all_img_paths) == 0:109 raise ValueError("Folder does not contain any images.")110 else:111 raise ValueError112 113 for input_img_path in all_img_paths:114 with Image.open(input_img_path) as image:115 if image.mode == "RGBA":116 image = image.convert("RGB")117 w, h = image.size118 119 if h % 64 != 0 or w % 64 != 0:120 width, height = map(lambda x: x - x % 64, (w, h))121 image = image.resize((width, height))122 print(123 f"WARNING: Your image is of size {h}x{w} which is not divisible by 64. We are resizing to {height}x{width}!"124 )125 126 image = ToTensor()(image)127 image = image * 2.0 - 1.0128 129 image = image.unsqueeze(0).to(device)130 H, W = image.shape[2:]131 assert image.shape[1] == 3132 F = 8133 C = 4134 shape = (num_frames, C, H // F, W // F)135 if (H, W) != (576, 1024):136 print(137 "WARNING: The conditioning frame you provided is not 576x1024. This leads to suboptimal performance as model was only trained on 576x1024. Consider increasing `cond_aug`."138 )139 if motion_bucket_id > 255:140 print(141 "WARNING: High motion bucket! This may lead to suboptimal performance."142 )143 144 if fps_id < 5:145 print("WARNING: Small fps value! This may lead to suboptimal performance.")146 147 if fps_id > 30:148 print("WARNING: Large fps value! This may lead to suboptimal performance.")149 150 value_dict = {}151 value_dict["motion_bucket_id"] = motion_bucket_id152 value_dict["fps_id"] = fps_id153 value_dict["cond_aug"] = cond_aug154 value_dict["cond_frames_without_noise"] = image155 value_dict["cond_frames"] = image + cond_aug * torch.randn_like(image)156 value_dict["cond_aug"] = cond_aug157 158 with torch.no_grad():159 with torch.autocast(device):160 batch, batch_uc = get_batch(161 get_unique_embedder_keys_from_conditioner(model.conditioner),162 value_dict,163 [1, num_frames],164 T=num_frames,165 device=device,166 )167 c, uc = model.conditioner.get_unconditional_conditioning(168 batch,169 batch_uc=batch_uc,170 force_uc_zero_embeddings=[171 "cond_frames",172 "cond_frames_without_noise",173 ],174 )175 176 for k in ["crossattn", "concat"]:177 uc[k] = repeat(uc[k], "b ... -> b t ...", t=num_frames)178 uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=num_frames)179 c[k] = repeat(c[k], "b ... -> b t ...", t=num_frames)180 c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=num_frames)181 182 randn = torch.randn(shape, device=device)183 184 additional_model_inputs = {}185 additional_model_inputs["image_only_indicator"] = torch.zeros(186 2, num_frames187 ).to(device)188 additional_model_inputs["num_video_frames"] = batch["num_video_frames"]189 190 def denoiser(input, sigma, c):191 return model.denoiser(192 model.model, input, sigma, c, **additional_model_inputs193 )194 195 samples_z = model.sampler(denoiser, randn, cond=c, uc=uc)196 model.en_and_decode_n_samples_a_time = decoding_t197 samples_x = model.decode_first_stage(samples_z)198 samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)199 200 os.makedirs(output_folder, exist_ok=True)201 base_count = len(glob(os.path.join(output_folder, "*.mp4")))202 video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")203 writer = cv2.VideoWriter(204 video_path,205 cv2.VideoWriter_fourcc(*"mp4v"),206 fps_id + 1,207 (samples.shape[-1], samples.shape[-2]),208 )209 210 samples = embed_watermark(samples)211 samples = filter(samples)212 vid = (213 (rearrange(samples, "t c h w -> t h w c") * 255)214 .cpu()215 .numpy()216 .astype(np.uint8)217 )218 for frame in vid:219 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)220 writer.write(frame)221 writer.release()222 223 return video_path, seed224 225def get_unique_embedder_keys_from_conditioner(conditioner):226 return list(set([x.input_key for x in conditioner.embedders]))227 228 229def get_batch(keys, value_dict, N, T, device):230 batch = {}231 batch_uc = {}232 233 for key in keys:234 if key == "fps_id":235 batch[key] = (236 torch.tensor([value_dict["fps_id"]])237 .to(device)238 .repeat(int(math.prod(N)))239 )240 elif key == "motion_bucket_id":241 batch[key] = (242 torch.tensor([value_dict["motion_bucket_id"]])243 .to(device)244 .repeat(int(math.prod(N)))245 )246 elif key == "cond_aug":247 batch[key] = repeat(248 torch.tensor([value_dict["cond_aug"]]).to(device),249 "1 -> b",250 b=math.prod(N),251 )252 elif key == "cond_frames":253 batch[key] = repeat(value_dict["cond_frames"], "1 ... -> b ...", b=N[0])254 elif key == "cond_frames_without_noise":255 batch[key] = repeat(256 value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]257 )258 else:259 batch[key] = value_dict[key]260 261 if T is not None:262 batch["num_video_frames"] = T263 264 for key in batch.keys():265 if key not in batch_uc and isinstance(batch[key], torch.Tensor):266 batch_uc[key] = torch.clone(batch[key])267 return batch, batch_uc268 269def resize_image(image_path, output_size=(1024, 576)):270 image = Image.open(image_path)271 # Calculate aspect ratios272 target_aspect = output_size[0] / output_size[1] # Aspect ratio of the desired size273 image_aspect = image.width / image.height # Aspect ratio of the original image274 275 # Resize then crop if the original image is larger276 if image_aspect > target_aspect:277 # Resize the image to match the target height, maintaining aspect ratio278 new_height = output_size[1]279 new_width = int(new_height * image_aspect)280 resized_image = image.resize((new_width, new_height), Image.LANCZOS)281 # Calculate coordinates for cropping282 left = (new_width - output_size[0]) / 2283 top = 0284 right = (new_width + output_size[0]) / 2285 bottom = output_size[1]286 else:287 # Resize the image to match the target width, maintaining aspect ratio288 new_width = output_size[0]289 new_height = int(new_width / image_aspect)290 resized_image = image.resize((new_width, new_height), Image.LANCZOS)291 # Calculate coordinates for cropping292 left = 0293 top = (new_height - output_size[1]) / 2294 right = output_size[0]295 bottom = (new_height + output_size[1]) / 2296 297 # Crop the image298 cropped_image = resized_image.crop((left, top, right, bottom))299 300 return cropped_image301 302with gr.Blocks() as demo:303 gr.Markdown('''# Community demo for Stable Video Diffusion - Img2Vid - XT ([model](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt), [paper](https://stability.ai/research/stable-video-diffusion-scaling-latent-video-diffusion-models-to-large-datasets))304#### Research release ([_non-commercial_](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt/blob/main/LICENSE)): generate `4s` vid from a single image at (`25 frames` at `6 fps`). Generation takes ~60s in an A100. [Join the waitlist for Stability's upcoming web experience](https://stability.ai/contact).305 ''')306 with gr.Row():307 with gr.Column():308 image = gr.Image(label="Upload your image", type="filepath")309 generate_btn = gr.Button("Generate")310 video = gr.Video()311 with gr.Accordion("Advanced options", open=False):312 seed = gr.Slider(label="Seed", value=42, randomize=True, minimum=0, maximum=max_64_bit_int, step=1)313 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)314 motion_bucket_id = gr.Slider(label="Motion bucket id", info="Controls how much motion to add/remove from the image", value=127, minimum=1, maximum=255)315 fps_id = gr.Slider(label="Frames per second", info="The length of your video in seconds will be 25/fps", value=6, minimum=5, maximum=30)316 317 image.upload(fn=resize_image, inputs=image, outputs=image, queue=False)318 generate_btn.click(fn=sample, inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id], outputs=[video, seed], api_name="video")319 320if __name__ == "__main__":321 demo.queue(max_size=20)322 demo.launch(share=True)