Spaces:
Build error
Build error
| from text_to_video.tuneavideo.pipelines.pipeline_text_to_video import TuneAVideoPipeline | |
| from text_to_video.tuneavideo.models.unet import UNet3DConditionModel | |
| import torch | |
| from diffusers import AutoencoderKL, DDIMScheduler | |
| from transformers import CLIPTextModel, CLIPTokenizer | |
| class TextToVideo(): | |
| def __init__(self,sd_path = None,motion_field_strength = 12, video_length = 8,t0 = 881, t1=941,use_cf_attn=True,use_motion_field=True) -> None: | |
| g = torch.Generator(device='cuda') | |
| g.manual_seed(22) | |
| self.g = g | |
| assert sd_path is not None | |
| print(f"Loading model SD-Net model file from {sd_path}") | |
| self.dtype = torch.float16 | |
| noise_scheduler = DDIMScheduler.from_pretrained( | |
| sd_path, subfolder="scheduler") | |
| tokenizer = CLIPTokenizer.from_pretrained( | |
| sd_path, subfolder="tokenizer") | |
| text_encoder = CLIPTextModel.from_pretrained( | |
| sd_path, subfolder="text_encoder") | |
| vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae") | |
| unet = UNet3DConditionModel.from_pretrained_2d( | |
| sd_path, subfolder="unet", use_cf_attn=use_cf_attn) | |
| self.pipe = TuneAVideoPipeline( | |
| vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, | |
| scheduler=DDIMScheduler.from_pretrained( | |
| sd_path, subfolder="scheduler") | |
| ).to('cuda').to(self.dtype) | |
| noise_scheduler.set_timesteps(50, device='cuda') | |
| # t0 parameter (DDIM backward from noise until t0) | |
| self.t0 = t0 | |
| # from t0 apply DDPM forward until t1 | |
| self.t1 = t1 | |
| self.use_foreground_motion_field = False # apply motion field on forground object (not used) | |
| # strength of motion field (delta_x = delta_y in Sect 3.3.1) | |
| self.motion_field_strength = motion_field_strength | |
| self.use_motion_field = use_motion_field # apply general motion field | |
| self.smooth_bg = False # temporally smooth background | |
| self.smooth_bg_strength = 0.4 # alpha = (1-self.smooth_bg_strength) in Eq (9) | |
| self.video_length = video_length | |
| def inference(self, prompt): | |
| prompt_compute = [prompt] | |
| xT = torch.randn((1, 4, 1, 64, 64), dtype=self.dtype, device="cuda") | |
| result = self.pipe(prompt_compute, | |
| video_length=self.video_length, | |
| height=512, | |
| width=512, | |
| num_inference_steps=50, | |
| guidance_scale=7.5, | |
| guidance_stop_step=1.0, | |
| t0=self.t0, | |
| t1=self.t1, | |
| xT=xT, | |
| use_foreground_motion_field=self.use_foreground_motion_field, | |
| motion_field_strength=self.motion_field_strength, | |
| use_motion_field=self.use_motion_field, | |
| smooth_bg=self.smooth_bg, | |
| smooth_bg_strength=self.smooth_bg_strength, | |
| generator=self.g) | |
| return result.videos[0] | |