Update app.py

app.py

@@ -84,13 +84,9 @@ class WaveCollapseTracker:
         self.snapshot = None
 
     def callback(self, pipe, step_index, timestep, callback_kwargs):
-        # 1. Trigger the step timer
         self.timer.step(step_index + 1)
-        
-        # 2. Extract current latents
         latents = callback_kwargs["latents"]
         
-        # 3. Wave Collapse Math
         if self.prev_latents is not None:
             delta = (latents - self.prev_latents).abs().mean(dim=1, keepdim=True)
             new_settled = delta < self.epsilon
@@ -159,6 +155,7 @@ def infer(
         try:
             torch.randn_like = _locked_randn_like
 
+            # We pass output_type="latent" so the pipeline stops before the VAE decode
             if auto_anti_prompt and prompt:
                 pos_hidden, pos_pooled = encode_prompt_sdxl(pipe, prompt, device)
                 neg_hidden = -pos_hidden
@@ -166,7 +163,7 @@ def infer(
 
                 wave_tracker.timer.start()
                 t_gen_start = time.time()
-                image = pipe(
+                pipeline_output = pipe(
                     prompt_embeds=pos_hidden,
                     negative_prompt_embeds=neg_hidden,
                     pooled_prompt_embeds=pos_pooled,
@@ -176,12 +173,13 @@ def infer(
                     width=width,
                     height=height,
                     generator=generator,
+                    output_type="latent",
                     callback_on_step_end=wave_tracker.callback,
-                ).images[0]
+                )
             else:
                 wave_tracker.timer.start()
                 t_gen_start = time.time()
-                image = pipe(
+                pipeline_output = pipe(
                     prompt=prompt,
                     negative_prompt=negative_prompt if negative_prompt else None,
                     guidance_scale=guidance_scale,
@@ -189,29 +187,32 @@ def infer(
                     width=width,
                     height=height,
                     generator=generator,
+                    output_type="latent",
                     callback_on_step_end=wave_tracker.callback,
-                ).images[0]
+                )
 
             t_gen_end = time.time()
+            final_latents = pipeline_output.images
             
-            # --- Decode the Wave Collapse Map ---
-            collapse_image = None
+            # --- Decode the Accumulated Wave Collapse Master Snapshot ---
             if wave_tracker.snapshot is not None and wave_tracker.cumulative_mask is not None:
-                unsettled = (~wave_tracker.cumulative_mask).expand_as(wave_tracker.prev_latents).to(wave_tracker.prev_latents.dtype)
-                final_snapshot = (wave_tracker.snapshot * (1.0 - unsettled)) + (wave_tracker.prev_latents * unsettled)
+                unsettled = (~wave_tracker.cumulative_mask).expand_as(final_latents).to(final_latents.dtype)
+                final_snapshot = (wave_tracker.snapshot * (1.0 - unsettled)) + (final_latents * unsettled)
+            else:
+                final_snapshot = final_latents
                 
-                with torch.no_grad():
-                    # Upcast to prevent VAE black-screen bug
-                    pipe.vae.to(dtype=torch.float32)
-                    final_snapshot = (final_snapshot / pipe.vae.config.scaling_factor).to(torch.float32)
-                    collapse_tensor = pipe.vae.decode(final_snapshot, return_dict=False)[0]
-                    collapse_image = pipe.image_processor.postprocess(collapse_tensor, output_type="pil")[0]
+            with torch.no_grad():
+                # Upcast to prevent VAE black-screen bug
+                pipe.vae.to(dtype=torch.float32)
+                final_snapshot_fp32 = (final_snapshot / pipe.vae.config.scaling_factor).to(torch.float32)
+                collapse_tensor = pipe.vae.decode(final_snapshot_fp32, return_dict=False)[0]
+                final_image = pipe.image_processor.postprocess(collapse_tensor, output_type="pil")[0]
 
             total_time = t_gen_end - t_start
             step_summary = timer.summary()
             status = f"{'CLIP Mirror ON' if auto_anti_prompt else 'Standard CFG'} | {step_summary} | Total: {total_time:.1f}s"
             
-            return image, collapse_image, seed, status
+            return final_image, seed, status
 
         finally:
             torch.randn_like = _original_randn_like
@@ -263,7 +264,6 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
 
         with gr.Column(scale=1):
             output_image = gr.Image(label="Final Generated Image")
-            collapse_map = gr.Image(label="Wave Collapse Visualization")
             
             with gr.Row():
                 output_seed = gr.Textbox(label="Used Seed", interactive=False)
@@ -276,7 +276,7 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
             width, height, guidance_scale, num_inference_steps,
             auto_anti_prompt, epsilon
         ],
-        outputs=[output_image, collapse_map, output_seed, status_display],
+        outputs=[output_image, output_seed, status_display],
     )
 
 if __name__ == "__main__":