sample

app.py

@@ -190,7 +190,7 @@ if __name__ == "__main__":
                         attn_map = gr.Image(show_label=False, show_download_button=False)
                     with gr.Tab(label="Segmentation Maps"):
                         gr.Markdown("### Character-level segmentation maps (using upscaled attention maps):")
-                        seg_map = gr.AnnotatedImage(height=384, show_label=False, show_download_button=False)
+                        seg_map = gr.AnnotatedImage(height=384, show_label=False)
 
         # examples
         examples = []

sgm/modules/diffusionmodules/sampling.py

@@ -268,12 +268,10 @@ class EulerEDMSampler(EDMSampler):
 
         return colormap
     
-    def save_segment_map(self, image, attn_maps, tokens=None, save_name=None):
+    def save_segment_map(self, H, W, attn_maps, tokens=None, save_name=None):
 
         colormap = self.create_pascal_label_colormap()
-        H, W = image.shape[-2:]
 
-        image_ = image*0.3
         sections = []
         for i in range(len(tokens)): 
             attn_map = attn_maps[i]
@@ -285,14 +283,11 @@ class EulerEDMSampler(EDMSampler):
             colored_attn_map = attn_map_t * color
             colored_attn_map = colored_attn_map.to(device=image_.device)
 
-            image_ += colored_attn_map*0.7
             sections.append(attn_map)
         
         section = np.stack(sections)
         np.save(f"./temp/seg_map/seg_{save_name}.npy", section)
 
-        save_image(image_, f"./temp/seg_map/seg_{save_name}.png", normalize=True)
-
     def get_init_noise(self, cfgs, model, cond, batch, uc=None):
 
         H, W = batch["target_size_as_tuple"][0]
@@ -375,8 +370,8 @@ class EulerEDMSampler(EDMSampler):
             local_loss = torch.zeros(1)
         if save_attn:
             attn_map = model.model.diffusion_model.save_attn_map(save_name=name, tokens=batch["label"][0])
-            denoised_decode = model.decode_first_stage(denoised) if denoised_decode is None else denoised_decode
-            self.save_segment_map(denoised_decode, attn_map, tokens=batch["label"][0], save_name=name)
+            H, W = batch["target_size_as_tuple"][0]
+            self.save_segment_map(H, W, attn_map, tokens=batch["label"][0], save_name=name)
 
         d = to_d(x, sigma_hat, denoised)
         dt = append_dims(next_sigma - sigma_hat, x.ndim)
@@ -557,8 +552,8 @@ class EulerEDMDualSampler(EulerEDMSampler):
             local_loss = torch.zeros(1)
         if save_attn:
             attn_map = model.model.diffusion_model.save_attn_map(save_name=name, save_single=True)
-            denoised_decode = model.decode_first_stage(denoised) if denoised_decode is None else denoised_decode
-            self.save_segment_map(denoised_decode, attn_map, tokens=batch["label"][0], save_name=name)
+            H, W = batch["target_size_as_tuple"][0]
+            self.save_segment_map(H, W, attn_map, tokens=batch["label"][0], save_name=name)
 
         d = to_d(x, sigma_hat, denoised)
         dt = append_dims(next_sigma - sigma_hat, x.ndim)