(args)
| 169 | return video_clip |
| 170 | |
| 171 | def main(args): |
| 172 | # Setup PyTorch: |
| 173 | if args.seed: |
| 174 | torch.manual_seed(args.seed) |
| 175 | torch.set_grad_enabled(False) |
| 176 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 177 | # device = "cpu" |
| 178 | |
| 179 | if args.ckpt is None: |
| 180 | raise ValueError("Please specify a checkpoint path using --ckpt <path>") |
| 181 | |
| 182 | # Load model: |
| 183 | latent_h = args.image_size[0] // 8 |
| 184 | latent_w = args.image_size[1] // 8 |
| 185 | args.image_h = args.image_size[0] |
| 186 | args.image_w = args.image_size[1] |
| 187 | args.latent_h = latent_h |
| 188 | args.latent_w = latent_w |
| 189 | print('loading model') |
| 190 | model = get_models(args).to(device) |
| 191 | |
| 192 | if args.enable_xformers_memory_efficient_attention: |
| 193 | if is_xformers_available(): |
| 194 | model.enable_xformers_memory_efficient_attention() |
| 195 | else: |
| 196 | raise ValueError("xformers is not available. Make sure it is installed correctly") |
| 197 | |
| 198 | # load model |
| 199 | ckpt_path = args.ckpt |
| 200 | state_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage)['ema'] |
| 201 | model.load_state_dict(state_dict) |
| 202 | print('loading succeed') |
| 203 | |
| 204 | model.eval() |
| 205 | pretrained_model_path = args.pretrained_model_path |
| 206 | diffusion = create_diffusion(str(args.num_sampling_steps)) |
| 207 | vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to(device) |
| 208 | text_encoder = TextEmbedder(pretrained_model_path).to(device) |
| 209 | if args.use_fp16: |
| 210 | print('Warnning: using half percision for inferencing!') |
| 211 | vae.to(dtype=torch.float16) |
| 212 | model.to(dtype=torch.float16) |
| 213 | text_encoder.to(dtype=torch.float16) |
| 214 | |
| 215 | # prompt: |
| 216 | prompt = args.text_prompt |
| 217 | if prompt ==[]: |
| 218 | prompt = args.input_path.split('/')[-1].split('.')[0].replace('_', ' ') |
| 219 | else: |
| 220 | prompt = prompt[0] |
| 221 | prompt_base = prompt.replace(' ','_') |
| 222 | prompt = prompt + args.additional_prompt |
| 223 | |
| 224 | if not os.path.exists(os.path.join(args.save_path)): |
| 225 | os.makedirs(os.path.join(args.save_path)) |
| 226 | video_input, researve_frames = get_input(args) # f,c,h,w |
| 227 | video_input = video_input.to(device).unsqueeze(0) # b,f,c,h,w |
| 228 | mask = mask_generation_before(args.mask_type, video_input.shape, video_input.dtype, device) # b,f,c,h,w |
no test coverage detected