Properly render cross attention as well.

app.py

@@ -45,10 +45,6 @@ def encode_image(image, prompt, concepts, seed, layer_start_index, noise_timeste
         if len(concepts) > 9:
             raise gr.Error("Please enter at most 9 concepts", duration=10)
 
-        print(f"Num samples: {num_samples}")
-        print(f"Layer start index: {layer_start_index}")
-        print(f"Noise timestep: {noise_timestep}")
-        print(image)
         image = image.convert("RGB")
 
         pipeline_output = pipeline.encode_image(

concept_attention/concept_attention_pipeline.py

@@ -354,8 +354,7 @@ class ConceptAttentionFluxPipeline():
         #     w=64
         # )
         concept_heatmaps = concept_heatmaps.to(torch.float32).detach().cpu().numpy()[0]
-        # cross_attention_maps = cross_attention_maps.to(torch.float32).detach().cpu().numpy()[0]
-        cross_attention_maps = concept_heatmaps
+        cross_attention_maps = cross_attention_maps.to(torch.float32).detach().cpu().numpy()[0]
         # Convert the torch heatmaps to PIL images.
         if return_pil_heatmaps:
             concept_heatmaps_min = concept_heatmaps.min()