Update app.py

app.py

@@ -158,11 +158,9 @@ def generate(
         all_timesteps_cpu = timesteps.cpu()
         timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 2)
         timesteps_split_for_state = [chunk for chunk in timesteps_split_np] # Store list of numpy arrays
-        
         # test with 2 segments
         segment_timesteps = torch.from_numpy(timesteps_split_np[0]).to("cuda")
         #segment_timesteps = timesteps
-        
         num_channels_latents = pipe.unet.config.in_channels
         latents = pipe.prepare_latents(
             batch_size=1, num_channels_latents=pipe.unet.config.in_channels, height=height, width=width,
@@ -189,21 +187,16 @@ def generate(
         state = torch.load(state_file, weights_only=False)
         latents = state["intermediate_latents"].to("cuda", dtype=torch.bfloat16)
         guidance_scale = state["guidance_scale"]
-        
+        seed = state["seed"]
         all_timesteps_cpu = state["all_timesteps"]
-        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 2)
-        timesteps_split_for_state = [chunk for chunk in timesteps_split_np] # Store list of numpy arrays
+        height = state["height"]
+        width = state["width"]
+        pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
         #timesteps_chunk_np = state["timesteps_split"][segment - 1]
-        segment_timesteps = torch.from_numpy(state["timesteps_split"][segment - 1]).to("cuda")
-        
+        #segment_timesteps = torch.from_numpy(state["timesteps_split"][segment - 1]).to("cuda")
             # test with 2 segments
         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)
         pooled_prompt_embeds = state["pooled_prompt_embeds"].to("cuda", dtype=torch.bfloat16)
@@ -215,6 +208,8 @@ def generate(
         add_time_ids = state["add_time_ids"].to("cuda", dtype=torch.bfloat16) # Original time IDs
         loop_add_time_ids = add_time_ids # Start with original loaded ones
         loop_add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
+        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 2)
+        timesteps_split_for_state = [chunk for chunk in timesteps_split_np] # Store list of numpy arrays
         added_cond_kwargs = {"text_embeds": unet_added_text_embeds, "time_ids": loop_add_time_ids}
         current_latents = latents # Start with loaded intermediate latents