Update app.py

app.py

@@ -151,13 +151,6 @@ 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)
-        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()
@@ -173,6 +166,12 @@ 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:
+        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)