Update injection_main.py

injection_main.py

@@ -236,15 +236,22 @@ def sample_disentangled(
     generative_latent *= pipe.scheduler.init_noise_sigma
 
     latents = start_latents.clone()
-
     latents = latents.repeat(len(prompt), 1, 1, 1)
     # randomly initialize the 1st latent for generation
-
     latents[1] = generative_latent
-    # assume that the first latent is used for reconstruction
+
+    num_intermediate_latents = len(intermediate_latents) if intermediate_latents is not None else 0
+
     for i in range(start_step, num_inference_steps):
-        if use_content_anchor:
-            latents[0] = intermediate_latents[-(i + 1)]
+        if use_content_anchor and intermediate_latents is not None:
+            # Ensure that the index is within bounds
+            if -i >= -num_intermediate_latents:
+                latents[0] = intermediate_latents[-i]
+            else:
+                # Handle case when the index is out of bounds
+                # You could use a default latent or skip this step
+                latents[0] = intermediate_latents[0]  # Example: use the first latent
+
         t = pipe.scheduler.timesteps[i]
 
         # Expand the latents if we are doing classifier free guidance