Update app.py

app.py

@@ -191,17 +191,17 @@ def generate(
         guidance_scale = state["guidance_scale"]
         all_timesteps_cpu = state["all_timesteps"]
         timesteps_split_for_state = state["timesteps_split"]
-        timesteps_chunk_np = state["timesteps_split"][segment - 1]
-        segment_timesteps = torch.from_numpy(timesteps_chunk_np).to("cuda")
+        #timesteps_chunk_np = state["timesteps_split"][segment - 1]
+        segment_timesteps = torch.from_numpy(timesteps_split_np[1]).to("cuda")
         seed = state["seed"]
         height = state["height"]
         width = state["width"]
         generator = torch.Generator(device='cuda').manual_seed(seed)
         pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
         prompt_embeds = state["prompt_embeds"].to("cuda", dtype=torch.bfloat16)
-        negative_prompt_embeds = state["negative_prompt_embeds"].to("cuda", dtype=torch.bfloat16) if state["negative_prompt_embeds"] is not None else None
+        negative_prompt_embeds = state["negative_prompt_embeds"].to("cuda", dtype=torch.bfloat16)
         pooled_prompt_embeds = state["pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16)
-        negative_pooled_prompt_embeds = state["negative_pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16) if state["negative_pooled_prompt_embeds"] is not None else None
+        negative_pooled_prompt_embeds = state["negative_pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16) if state["negative_pooled_prompt_embeds"]
         unet_prompt_embeds = prompt_embeds
         unet_added_text_embeds = pooled_prompt_embeds
         unet_prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
@@ -222,7 +222,7 @@ def generate(
         current_latents = pipe.scheduler.step(noise_pred, t, current_latents, generator=generator, return_dict=False)[0]
         
     intermediate_latents_cpu = current_latents.detach().cpu() # Latents after first half, moved to CPU
-        
+
     # test with 2 segments
     if segment==2:
         final_latents = current_latents
@@ -243,7 +243,7 @@ def generate(
         original_negative_prompt_embeds_cpu = negative_prompt_embeds.cpu()
         original_pooled_prompt_embeds_cpu = pooled_prompt_embeds.cpu()
         original_negative_pooled_prompt_embeds_cpu = negative_pooled_prompt_embeds.cpu()
-        original_add_time_ids_cpu = add_time_ids.cpu() # Saved before CFG duplication
+        original_add_time_ids_cpu = add_time_ids.cpu()
         state = {
         "intermediate_latents": intermediate_latents_cpu,
         "all_timesteps": all_timesteps_cpu, # Save full list generated by scheduler