Joeythemonster/Text2Video-Zero
0
1from enum import Enum2import gc3import numpy as np4import tomesd5import torch6 7from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionControlNetPipeline, ControlNetModel, UNet2DConditionModel8from diffusers.schedulers import EulerAncestralDiscreteScheduler, DDIMScheduler9from text_to_video_pipeline import TextToVideoPipeline10 11import utils12import gradio_utils13import os14on_huggingspace = os.environ.get("SPACE_AUTHOR_NAME") == "PAIR"15 16from einops import rearrange17 18 19class ModelType(Enum):20 Pix2Pix_Video = 1,21 Text2Video = 2,22 ControlNetCanny = 3,23 ControlNetCannyDB = 4,24 ControlNetPose = 5,25 ControlNetDepth = 6,26 27 28class Model:29 def __init__(self, device, dtype, **kwargs):30 self.device = device31 self.dtype = dtype32 self.generator = torch.Generator(device=device)33 self.pipe_dict = {34 ModelType.Pix2Pix_Video: StableDiffusionInstructPix2PixPipeline,35 ModelType.Text2Video: TextToVideoPipeline,36 ModelType.ControlNetCanny: StableDiffusionControlNetPipeline,37 ModelType.ControlNetCannyDB: StableDiffusionControlNetPipeline,38 ModelType.ControlNetPose: StableDiffusionControlNetPipeline,39 ModelType.ControlNetDepth: StableDiffusionControlNetPipeline,40 }41 self.controlnet_attn_proc = utils.CrossFrameAttnProcessor(42 unet_chunk_size=2)43 self.pix2pix_attn_proc = utils.CrossFrameAttnProcessor(44 unet_chunk_size=3)45 self.text2video_attn_proc = utils.CrossFrameAttnProcessor(46 unet_chunk_size=2)47 48 self.pipe = None49 self.model_type = None50 51 self.states = {}52 self.model_name = ""53 54 def set_model(self, model_type: ModelType, model_id: str, **kwargs):55 if hasattr(self, "pipe") and self.pipe is not None:56 del self.pipe57 torch.cuda.empty_cache()58 gc.collect()59 safety_checker = kwargs.pop('safety_checker', None)60 self.pipe = self.pipe_dict[model_type].from_pretrained(61 model_id, safety_checker=safety_checker, **kwargs).to(self.device).to(self.dtype)62 self.model_type = model_type63 self.model_name = model_id64 65 def inference_chunk(self, frame_ids, **kwargs):66 if not hasattr(self, "pipe") or self.pipe is None:67 return68 69 prompt = np.array(kwargs.pop('prompt'))70 negative_prompt = np.array(kwargs.pop('negative_prompt', ''))71 latents = None72 if 'latents' in kwargs:73 latents = kwargs.pop('latents')[frame_ids]74 if 'image' in kwargs:75 kwargs['image'] = kwargs['image'][frame_ids]76 if 'video_length' in kwargs:77 kwargs['video_length'] = len(frame_ids)78 if self.model_type == ModelType.Text2Video:79 kwargs["frame_ids"] = frame_ids80 return self.pipe(prompt=prompt[frame_ids].tolist(),81 negative_prompt=negative_prompt[frame_ids].tolist(),82 latents=latents,83 generator=self.generator,84 **kwargs)85 86 def inference(self, split_to_chunks=False, chunk_size=2, **kwargs):87 if not hasattr(self, "pipe") or self.pipe is None:88 return89 90 if "merging_ratio" in kwargs:91 merging_ratio = kwargs.pop("merging_ratio")92 93 # if merging_ratio > 0:94 tomesd.apply_patch(self.pipe, ratio=merging_ratio)95 seed = kwargs.pop('seed', 0)96 if seed < 0:97 seed = self.generator.seed()98 kwargs.pop('generator', '')99 100 if 'image' in kwargs:101 f = kwargs['image'].shape[0]102 else:103 f = kwargs['video_length']104 105 assert 'prompt' in kwargs106 prompt = [kwargs.pop('prompt')] * f107 negative_prompt = [kwargs.pop('negative_prompt', '')] * f108 109 frames_counter = 0110 111 # Processing chunk-by-chunk112 if split_to_chunks:113 chunk_ids = np.arange(0, f, chunk_size - 1)114 result = []115 for i in range(len(chunk_ids)):116 ch_start = chunk_ids[i]117 ch_end = f if i == len(chunk_ids) - 1 else chunk_ids[i + 1]118 frame_ids = [0] + list(range(ch_start, ch_end))119 self.generator.manual_seed(seed)120 print(f'Processing chunk {i + 1} / {len(chunk_ids)}')121 result.append(self.inference_chunk(frame_ids=frame_ids,122 prompt=prompt,123 negative_prompt=negative_prompt,124 **kwargs).images[1:])125 frames_counter += len(chunk_ids)-1126 if on_huggingspace and frames_counter >= 80:127 break128 result = np.concatenate(result)129 return result130 else:131 self.generator.manual_seed(seed)132 return self.pipe(prompt=prompt, negative_prompt=negative_prompt, generator=self.generator, **kwargs).images133 134 def process_controlnet_canny(self,135 video_path,136 prompt,137 chunk_size=2,138 watermark='Picsart AI Research',139 merging_ratio=0.0,140 num_inference_steps=20,141 controlnet_conditioning_scale=1.0,142 guidance_scale=9.0,143 seed=42,144 eta=0.0,145 low_threshold=100,146 high_threshold=200,147 resolution=512,148 use_cf_attn=True,149 save_path=None):150 print("Module Canny")151 video_path = gradio_utils.edge_path_to_video_path(video_path)152 if self.model_type != ModelType.ControlNetCanny:153 controlnet = ControlNetModel.from_pretrained(154 "lllyasviel/sd-controlnet-canny")155 self.set_model(ModelType.ControlNetCanny,156 model_id="runwayml/stable-diffusion-v1-5", controlnet=controlnet)157 158 self.pipe.scheduler = DDIMScheduler.from_config(159 self.pipe.scheduler.config)160 if use_cf_attn:161 self.pipe.unet.set_attn_processor(162 processor=self.controlnet_attn_proc)163 self.pipe.controlnet.set_attn_processor(164 processor=self.controlnet_attn_proc)165 166 added_prompt = 'best quality, extremely detailed'167 negative_prompts = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'168 169 video, fps = utils.prepare_video(170 video_path, resolution, self.device, self.dtype, False)171 control = utils.pre_process_canny(172 video, low_threshold, high_threshold).to(self.device).to(self.dtype)173 174 # canny_to_save = list(rearrange(control, 'f c w h -> f w h c').cpu().detach().numpy())175 # _ = utils.create_video(canny_to_save, 4, path="ddxk.mp4", watermark=None)176 177 f, _, h, w = video.shape178 self.generator.manual_seed(seed)179 latents = torch.randn((1, 4, h//8, w//8), dtype=self.dtype,180 device=self.device, generator=self.generator)181 latents = latents.repeat(f, 1, 1, 1)182 result = self.inference(image=control,183 prompt=prompt + ', ' + added_prompt,184 height=h,185 width=w,186 negative_prompt=negative_prompts,187 num_inference_steps=num_inference_steps,188 guidance_scale=guidance_scale,189 controlnet_conditioning_scale=controlnet_conditioning_scale,190 eta=eta,191 latents=latents,192 seed=seed,193 output_type='numpy',194 split_to_chunks=True,195 chunk_size=chunk_size,196 merging_ratio=merging_ratio,197 )198 return utils.create_video(result, fps, path=save_path, watermark=gradio_utils.logo_name_to_path(watermark))199 200 def process_controlnet_depth(self,201 video_path,202 prompt,203 chunk_size=2,204 watermark='Picsart AI Research',205 merging_ratio=0.0,206 num_inference_steps=20,207 controlnet_conditioning_scale=1.0,208 guidance_scale=9.0,209 seed=42,210 eta=0.0,211 resolution=512,212 use_cf_attn=True,213 save_path=None):214 print("Module Depth")215 video_path = gradio_utils.depth_path_to_video_path(video_path)216 if self.model_type != ModelType.ControlNetDepth:217 controlnet = ControlNetModel.from_pretrained(218 "lllyasviel/sd-controlnet-depth")219 self.set_model(ModelType.ControlNetDepth,220 model_id="runwayml/stable-diffusion-v1-5", controlnet=controlnet)221 self.pipe.scheduler = DDIMScheduler.from_config(222 self.pipe.scheduler.config)223 if use_cf_attn:224 self.pipe.unet.set_attn_processor(225 processor=self.controlnet_attn_proc)226 self.pipe.controlnet.set_attn_processor(227 processor=self.controlnet_attn_proc)228 229 added_prompt = 'best quality, extremely detailed'230 negative_prompts = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'231 232 video, fps = utils.prepare_video(233 video_path, resolution, self.device, self.dtype, False)234 control = utils.pre_process_depth(235 video).to(self.device).to(self.dtype)236 237 # depth_map_to_save = list(rearrange(control, 'f c w h -> f w h c').cpu().detach().numpy())238 # _ = utils.create_video(depth_map_to_save, 4, path="ddxk.mp4", watermark=None)239 240 f, _, h, w = video.shape241 self.generator.manual_seed(seed)242 latents = torch.randn((1, 4, h//8, w//8), dtype=self.dtype,243 device=self.device, generator=self.generator)244 latents = latents.repeat(f, 1, 1, 1)245 result = self.inference(image=control,246 prompt=prompt + ', ' + added_prompt,247 height=h,248 width=w,249 negative_prompt=negative_prompts,250 num_inference_steps=num_inference_steps,251 guidance_scale=guidance_scale,252 controlnet_conditioning_scale=controlnet_conditioning_scale,253 eta=eta,254 latents=latents,255 seed=seed,256 output_type='numpy',257 split_to_chunks=True,258 chunk_size=chunk_size,259 merging_ratio=merging_ratio,260 )261 return utils.create_video(result, fps, path=save_path, watermark=gradio_utils.logo_name_to_path(watermark))262 263 def process_controlnet_pose(self,264 video_path,265 prompt,266 chunk_size=2,267 watermark='Picsart AI Research',268 merging_ratio=0.0,269 num_inference_steps=20,270 controlnet_conditioning_scale=1.0,271 guidance_scale=9.0,272 seed=42,273 eta=0.0,274 resolution=512,275 use_cf_attn=True,276 save_path=None):277 print("Module Pose")278 video_path = gradio_utils.motion_to_video_path(video_path)279 if self.model_type != ModelType.ControlNetPose:280 controlnet = ControlNetModel.from_pretrained(281 "fusing/stable-diffusion-v1-5-controlnet-openpose")282 self.set_model(ModelType.ControlNetPose,283 model_id="runwayml/stable-diffusion-v1-5", controlnet=controlnet)284 self.pipe.scheduler = DDIMScheduler.from_config(285 self.pipe.scheduler.config)286 if use_cf_attn:287 self.pipe.unet.set_attn_processor(288 processor=self.controlnet_attn_proc)289 self.pipe.controlnet.set_attn_processor(290 processor=self.controlnet_attn_proc)291 292 video_path = gradio_utils.motion_to_video_path(293 video_path) if 'Motion' in video_path else video_path294 295 added_prompt = 'best quality, extremely detailed, HD, ultra-realistic, 8K, HQ, masterpiece, trending on artstation, art, smooth'296 negative_prompts = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly, unrealistic'297 298 video, fps = utils.prepare_video(299 video_path, resolution, self.device, self.dtype, False, output_fps=4)300 control = utils.pre_process_pose(301 video, apply_pose_detect=False).to(self.device).to(self.dtype)302 f, _, h, w = video.shape303 self.generator.manual_seed(seed)304 latents = torch.randn((1, 4, h//8, w//8), dtype=self.dtype,305 device=self.device, generator=self.generator)306 latents = latents.repeat(f, 1, 1, 1)307 result = self.inference(image=control,308 prompt=prompt + ', ' + added_prompt,309 height=h,310 width=w,311 negative_prompt=negative_prompts,312 num_inference_steps=num_inference_steps,313 guidance_scale=guidance_scale,314 controlnet_conditioning_scale=controlnet_conditioning_scale,315 eta=eta,316 latents=latents,317 seed=seed,318 output_type='numpy',319 split_to_chunks=True,320 chunk_size=chunk_size,321 merging_ratio=merging_ratio,322 )323 return utils.create_gif(result, fps, path=save_path, watermark=gradio_utils.logo_name_to_path(watermark))324 325 def process_controlnet_canny_db(self,326 db_path,327 video_path,328 prompt,329 chunk_size=2,330 watermark='Picsart AI Research',331 merging_ratio=0.0,332 num_inference_steps=20,333 controlnet_conditioning_scale=1.0,334 guidance_scale=9.0,335 seed=42,336 eta=0.0,337 low_threshold=100,338 high_threshold=200,339 resolution=512,340 use_cf_attn=True,341 save_path=None):342 print("Module Canny_DB")343 db_path = gradio_utils.get_model_from_db_selection(db_path)344 video_path = gradio_utils.get_video_from_canny_selection(video_path)345 # Load db and controlnet weights346 if 'db_path' not in self.states or db_path != self.states['db_path']:347 controlnet = ControlNetModel.from_pretrained(348 "lllyasviel/sd-controlnet-canny")349 self.set_model(ModelType.ControlNetCannyDB,350 model_id=db_path, controlnet=controlnet)351 self.pipe.scheduler = DDIMScheduler.from_config(352 self.pipe.scheduler.config)353 self.states['db_path'] = db_path354 355 if use_cf_attn:356 self.pipe.unet.set_attn_processor(357 processor=self.controlnet_attn_proc)358 self.pipe.controlnet.set_attn_processor(359 processor=self.controlnet_attn_proc)360 361 added_prompt = 'best quality, extremely detailed'362 negative_prompts = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'363 364 video, fps = utils.prepare_video(365 video_path, resolution, self.device, self.dtype, False)366 control = utils.pre_process_canny(367 video, low_threshold, high_threshold).to(self.device).to(self.dtype)368 f, _, h, w = video.shape369 self.generator.manual_seed(seed)370 latents = torch.randn((1, 4, h//8, w//8), dtype=self.dtype,371 device=self.device, generator=self.generator)372 latents = latents.repeat(f, 1, 1, 1)373 result = self.inference(image=control,374 prompt=prompt + ', ' + added_prompt,375 height=h,376 width=w,377 negative_prompt=negative_prompts,378 num_inference_steps=num_inference_steps,379 guidance_scale=guidance_scale,380 controlnet_conditioning_scale=controlnet_conditioning_scale,381 eta=eta,382 latents=latents,383 seed=seed,384 output_type='numpy',385 split_to_chunks=True,386 chunk_size=chunk_size,387 merging_ratio=merging_ratio,388 )389 return utils.create_gif(result, fps, path=save_path, watermark=gradio_utils.logo_name_to_path(watermark))390 391 def process_pix2pix(self,392 video,393 prompt,394 resolution=512,395 seed=0,396 image_guidance_scale=1.0,397 start_t=0,398 end_t=-1,399 out_fps=-1,400 chunk_size=2,401 watermark='Picsart AI Research',402 merging_ratio=0.0,403 use_cf_attn=True,404 save_path=None,):405 print("Module Pix2Pix")406 if self.model_type != ModelType.Pix2Pix_Video:407 self.set_model(ModelType.Pix2Pix_Video,408 model_id="timbrooks/instruct-pix2pix")409 self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(410 self.pipe.scheduler.config)411 if use_cf_attn:412 self.pipe.unet.set_attn_processor(413 processor=self.pix2pix_attn_proc)414 video, fps = utils.prepare_video(415 video, resolution, self.device, self.dtype, True, start_t, end_t, out_fps)416 self.generator.manual_seed(seed)417 result = self.inference(image=video,418 prompt=prompt,419 seed=seed,420 output_type='numpy',421 num_inference_steps=50,422 image_guidance_scale=image_guidance_scale,423 split_to_chunks=True,424 chunk_size=chunk_size,425 merging_ratio=merging_ratio426 )427 return utils.create_video(result, fps, path=save_path, watermark=gradio_utils.logo_name_to_path(watermark))428 429 def process_text2video(self,430 prompt,431 model_name="dreamlike-art/dreamlike-photoreal-2.0",432 motion_field_strength_x=12,433 motion_field_strength_y=12,434 t0=44,435 t1=47,436 n_prompt="",437 chunk_size=2,438 video_length=8,439 watermark='Picsart AI Research',440 merging_ratio=0.0,441 seed=0,442 resolution=512,443 fps=2,444 use_cf_attn=True,445 use_motion_field=True,446 smooth_bg=False,447 smooth_bg_strength=0.4,448 path=None):449 print("Module Text2Video")450 if self.model_type != ModelType.Text2Video or model_name != self.model_name:451 print("Model update")452 unet = UNet2DConditionModel.from_pretrained(453 model_name, subfolder="unet")454 self.set_model(ModelType.Text2Video,455 model_id=model_name, unet=unet)456 self.pipe.scheduler = DDIMScheduler.from_config(457 self.pipe.scheduler.config)458 if use_cf_attn:459 self.pipe.unet.set_attn_processor(460 processor=self.text2video_attn_proc)461 self.generator.manual_seed(seed)462 463 added_prompt = "high quality, HD, 8K, trending on artstation, high focus, dramatic lighting"464 negative_prompts = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly, unrealistic'465 466 prompt = prompt.rstrip()467 if len(prompt) > 0 and (prompt[-1] == "," or prompt[-1] == "."):468 prompt = prompt.rstrip()[:-1]469 prompt = prompt.rstrip()470 prompt = prompt + ", "+added_prompt471 if len(n_prompt) > 0:472 negative_prompt = n_prompt473 else:474 negative_prompt = None475 476 result = self.inference(prompt=prompt,477 video_length=video_length,478 height=resolution,479 width=resolution,480 num_inference_steps=50,481 guidance_scale=7.5,482 guidance_stop_step=1.0,483 t0=t0,484 t1=t1,485 motion_field_strength_x=motion_field_strength_x,486 motion_field_strength_y=motion_field_strength_y,487 use_motion_field=use_motion_field,488 smooth_bg=smooth_bg,489 smooth_bg_strength=smooth_bg_strength,490 seed=seed,491 output_type='numpy',492 negative_prompt=negative_prompt,493 merging_ratio=merging_ratio,494 split_to_chunks=True,495 chunk_size=chunk_size,496 )497 return utils.create_video(result, fps, path=path, watermark=gradio_utils.logo_name_to_path(watermark))498