Update app.py

app.py

@@ -151,6 +151,10 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         size = state["width"]
         pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
         timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
+        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)
     if segment==9:
         pipe.trasformer.to('cpu')
         torch.cuda.empty_cache()
@@ -167,10 +171,6 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         return video_out_file, seed
     else:
         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)