Update app.py

app.py

@@ -138,34 +138,21 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         torch.cuda.empty_cache()
         guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
     else:
-        pipe.vae.to("cpu")
-        torch.cuda.empty_cache()
-        transformer_dtype = pipe.transformer.dtype
         state_file = f"SkyReel_{segment-1}_{seed}.pt"
+        transformer_dtype = pipe.transformer.dtype
         state = torch.load(state_file, weights_only=False)
         generator = torch.Generator(device='cuda').manual_seed(seed)
         latents = state["intermediate_latents"].to("cuda", dtype=torch.bfloat16)
-        guidance_scale = state["guidance_scale"]
         all_timesteps_cpu = state["all_timesteps"]
-        size = state["height"]
-        size = state["width"]
         pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
-        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
-        segment_timesteps = torch.from_numpy(timesteps_split_np[segment - 1]).to("cuda")
-        prompt_embeds = state["prompt_embeds"].to("cuda", dtype=torch.bfloat16)
-        pooled_prompt_embeds = state["pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16)
-        prompt_attention_mask = state["prompt_attention_mask"].to("cuda", dtype=torch.bfloat16)
-        image_latents = state["image_latents"].to("cuda", dtype=torch.bfloat16)
-        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
-    #pipe.transformer.to("cuda")
     if segment==9:
         pipe.trasformer.to('cpu')
         torch.cuda.empty_cache()
         pipe.vae.to("cuda")
         latents = latents.to(pipe.vae.dtype) / pipe.vae.config.scaling_factor
-        #with torch.no_grad():
-        video = pipe.vae.decode(latents, return_dict=False)[0]
-        video = pipe.video_processor.postprocess_video(video)
+        with torch.no_grad():
+            video = pipe.vae.decode(latents, return_dict=False)[0]
+            video = pipe.video_processor.postprocess_video(video)
         # return HunyuanVideoPipelineOutput(frames=video)
         save_dir = f"./"
         video_out_file = f"{save_dir}/{seed}.mp4"
@@ -173,6 +160,18 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         export_to_video(output, video_out_file, fps=24)
         return video_out_file, seed
     else:
+        pipe.vae.to("cpu")
+        torch.cuda.empty_cache()
+        guidance_scale = state["guidance_scale"]
+        size = state["height"]
+        size = state["width"]
+        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
+        segment_timesteps = torch.from_numpy(timesteps_split_np[segment - 1]).to("cuda")
+        prompt_embeds = state["prompt_embeds"].to("cuda", dtype=torch.bfloat16)
+        pooled_prompt_embeds = state["pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16)
+        prompt_attention_mask = state["prompt_attention_mask"].to("cuda", dtype=torch.bfloat16)
+        image_latents = state["image_latents"].to("cuda", dtype=torch.bfloat16)
+        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
         for i, t in enumerate(pipe.progress_bar(segment_timesteps)):
                 latents = latents.to(transformer_dtype)
                 latent_model_input = torch.cat([latents] * 2)