Update app.py

app.py

@@ -154,9 +154,10 @@ def generate(
         )
         pipe.scheduler.set_timesteps(num_inference_steps, device=torch.device('cuda'))
         timesteps = pipe.scheduler.timesteps
-        
+        all_timesteps_cpu = timesteps.cpu()
+        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
+        timesteps_split_for_state = [chunk for chunk in timesteps_split_np] # Store list of numpy arrays
         split_point = num_inference_steps // 8
-        
         segment_timesteps = timesteps[:split_point].to("cuda") # Slice from CPU list, move to GPU
         num_channels_latents = pipe.unet.config.in_channels
         latents = pipe.prepare_latents(
@@ -192,7 +193,7 @@ def generate(
         guidance_scale = state["guidance_scale"]
         all_timesteps_cpu = state["all_timesteps"]
         
-        timesteps_split = state["timesteps_split"]
+        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")
         
@@ -218,7 +219,6 @@ def generate(
 
         added_cond_kwargs = {"text_embeds": unet_added_text_embeds, "time_ids": loop_add_time_ids}
         current_latents = latents # Start with loaded intermediate latents
-
     
     for i, t in enumerate(pipe.progress_bar(segment_timesteps)): # Only first half timesteps
         latent_model_input = torch.cat([current_latents] * 2) if guidance_scale > 1.0 else current_latents
@@ -255,9 +255,7 @@ def generate(
         "negative_pooled_prompt_embeds": original_negative_pooled_prompt_embeds_cpu,
         "add_time_ids": original_add_time_ids_cpu, # Save ORIGINAL time IDs
         "guidance_scale": guidance_scale,
-            
-        "timesteps_split": timesteps_split,
-            
+        "timesteps_split": timesteps_split_for_state,
         "seed": seed,
         "prompt": prompt, # Save originals for reference/verification
         "negative_prompt": negative_prompt,
@@ -293,7 +291,7 @@ with gr.Blocks(theme=gr.themes.Origin()) as demo:
             slider = gr.Slider(
                 minimum=1,
                 maximum=250,
-                value=[i * 125],
+                value=[i * (master_slider // 8)],
                 step=1,
                 label=f"Range {i + 1}",
             )