Update app.py

app.py

@@ -121,7 +121,7 @@ def main():
                         out2 = pipe(new_prompt, num_inference_steps=inf_steps, generator=gen)
                         image2 = np.array(out2.images[0]) / 255
             else:
-                with trace(pipe) as tc:
+                with trace(pipe, load_heads=False, save_heads=False) as tc:
                     out = pipe(prompt, num_inference_steps=inf_steps, generator=gen)
                     image = np.array(out.images[0]) / 255
                     heat_map = tc.compute_global_heat_map()