Update app.py

app.py

@@ -158,8 +158,11 @@ def generate(
         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 = torch.from_numpy(timesteps_split_np[0]).to("cuda")
+        
+        # test with one segment
+        #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,
@@ -209,8 +212,8 @@ 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
+    for i, t in enumerate(pipe.progress_bar(segment_timesteps)):
+        latent_model_input = torch.cat([current_latents] * 2)
         latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
         with torch.no_grad():
           noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=unet_prompt_embeds,added_cond_kwargs=added_cond_kwargs, return_dict=False)[0]
@@ -219,15 +222,16 @@ 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
-
-    if segment==8:
-        final_latents = current_latents # Latents after the second half loop
+        
+    # test with one segment
+    if segment==1:
+        final_latents = current_latents
         final_latents = final_latents / pipe.vae.config.scaling_factor
         with torch.no_grad():
-            image = pipe.vae.decode(final_latents, return_dict=False)[0] # VAE might prefer fp16/fp32
+            image = pipe.vae.decode(final_latents, return_dict=False)[0]
         image = pipe.image_processor.postprocess(image, output_type="pil")[0]
         output_image_file = f"rv_L_{seed}.png"
-        image.save(output_image_file) # Use output_image_file defined earlier
+        image.save(output_image_file)
         #timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         #upload_to_ftp(filename) 
         #uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)