examples

app.py

@@ -171,7 +171,7 @@ if __name__ == "__main__":
     model = init_model(cfgs)
     sampler = init_sampling(cfgs)
     global_index = 0
-    resize = Resize((cfgs.H, cfgs.W))
+    resize = Resize((cfgs.H, cfgs.W), antialias=True)
 
     block = gr.Blocks().queue()
     with block:
@@ -202,7 +202,7 @@ if __name__ == "__main__":
             with gr.Column():
 
                 input_blk = gr.Image(source='upload', tool='sketch', type="numpy", label="Input", height=512)
-                gr.Markdown("Notice: please draw horizontally to indicate only **one** masked area.")
+                gr.Markdown("Notice: please draw horizontally to indicate only **one** masked area. The image may be cropped automatically into a proper scale.")
                 text = gr.Textbox(label="Text to render: (1~12 characters)", info="the text you want to render at the masked region")
                 run_button = gr.Button(variant="primary")
 
@@ -210,9 +210,9 @@ if __name__ == "__main__":
 
                     num_samples = gr.Slider(label="Images", info="number of generated images, locked as 1", minimum=1, maximum=1, value=1, step=1)
                     steps = gr.Slider(label="Steps", info ="denoising sampling steps", minimum=1, maximum=200, value=50, step=1)
-                    scale = gr.Slider(label="Guidance Scale", info="the scale of classifier-free guidance (CFG)", minimum=0.0, maximum=10.0, value=5.0, step=0.1)
+                    scale = gr.Slider(label="Guidance Scale", info="the scale of classifier-free guidance (CFG)", minimum=0.0, maximum=10.0, value=4.0, step=0.1)
                     seed = gr.Slider(label="Seed", info="random seed for noise initialization", minimum=0, maximum=2147483647, step=1, randomize=True)
-                    show_detail = gr.Checkbox(label="Show Detail", info="show the additional visualization results", value=False)
+                    show_detail = gr.Checkbox(label="Show Detail", info="show the additional visualization results", value=True)
 
             with gr.Column():
 

sgm/modules/diffusionmodules/guiders.py

@@ -13,6 +13,7 @@ class VanillaCFG:
     def __init__(self, scale, dyn_thresh_config=None):
         scale_schedule = lambda scale, sigma: scale  # independent of step
         self.scale_schedule = partial(scale_schedule, scale)
+        self.scale_value = scale
         self.dyn_thresh = instantiate_from_config(
             default(
                 dyn_thresh_config,
@@ -24,7 +25,7 @@ class VanillaCFG:
 
     def __call__(self, x, sigma):
         x_u, x_c = x.chunk(2)
-        scale_value = self.scale_schedule(sigma)
+        scale_value = self.scale_value
         x_pred = self.dyn_thresh(x_u, x_c, scale_value)
         return x_pred
 

sgm/modules/diffusionmodules/sampling.py

@@ -7,7 +7,6 @@ from typing import Dict, Union
 
 import imageio
 import torch
-import json
 import numpy as np
 import torch.nn.functional as F
 from omegaconf import ListConfig, OmegaConf
@@ -252,47 +251,15 @@ class EulerEDMSampler(EDMSampler):
 
         return x
 
-    def create_pascal_label_colormap(self):
-        """
-        PASCAL VOC 分割数据集的类别标签颜色映射label colormap
+    def save_segment_map(self, attn_maps, tokens=None, save_name=None):
 
-        返回:
-            可视化分割结果的颜色映射Colormap
-        """
-        colormap = np.zeros((256, 3), dtype=int)
-        ind = np.arange(256, dtype=int)
-
-        for shift in reversed(range(8)):
-            for channel in range(3):
-                colormap[:, channel] |= ((ind >> channel) & 1) << shift
-            ind >>= 3
-
-        return colormap
-    
-    def save_segment_map(self, image, 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]
-            attn_map_t = np.tile(attn_map[None], (1,3,1,1)) # b, 3, h, w
-            attn_map_t = torch.from_numpy(attn_map_t)
-            attn_map_t = F.interpolate(attn_map_t, (W, H))
-
-            color = torch.from_numpy(colormap[i+1][None,:,None,None] / 255.0)
-            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)
+        np.save(f"./temp/seg_map/seg_{save_name}.npy", section)
 
     def get_init_noise(self, cfgs, model, cond, batch, uc=None):
 
@@ -376,8 +343,7 @@ 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)
+            self.save_segment_map(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)
@@ -410,7 +376,7 @@ class EulerEDMSampler(EDMSampler):
 
             alpha = 20 * np.sqrt(scales[i])
             update = aae_enabled
-            save_loss = detailed
+            save_loss = aae_enabled
             save_attn = detailed and (i == (num_sigmas-1)//2)
             save_inter = aae_enabled
 
@@ -452,7 +418,7 @@ class EulerEDMSampler(EDMSampler):
             imageio.mimsave(f"./temp/inters/{name}.gif", inters, 'GIF', duration=0.02)
 
         return x
-    
+
 
 class EulerEDMDualSampler(EulerEDMSampler):
 
@@ -557,9 +523,8 @@ class EulerEDMDualSampler(EulerEDMSampler):
         else:
             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)
+            attn_map = model.model.diffusion_model.save_attn_map(save_name=name, tokens=batch["label"][0])
+            self.save_segment_map(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)
@@ -632,7 +597,7 @@ class EulerEDMDualSampler(EulerEDMSampler):
         print(f"Local losses: {local_losses}")
 
         if len(inters) > 0:
-            imageio.mimsave(f"./temp/inters/{name}.gif", inters, 'GIF', duration=0.1)
+            imageio.mimsave(f"./temp/inters/{name}.gif", inters, 'GIF', duration=0.02)
 
         return x
 

temp/seg_map/seg_1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fed32518482697a99ffa93f572123251ff1a1ce344e6c87108c8e6e5344428cb
+size 32896

temp/seg_map/seg_10.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63ad403c856b5eb96cb1512d84fc5b030ea3fc0a043d3a80d3eeb37472cf343e
+size 24704

temp/seg_map/seg_11.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2181e1e25621b1dcb55d2e59df2b168d786021936a99e6e954f0825474414714
+size 24704

temp/seg_map/seg_12.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3ddb1524d276a3c1bdb101f9d6239adde117f441205ff44c325e96f11c2cef6
+size 20608

temp/seg_map/seg_13.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2f93300e573de546ae8ef84db4f49010b31df98561bd658f8da017cf277cad3
+size 24704

temp/seg_map/seg_14.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bc4905adb15d7bf5b7b902f7ed553bbf465e0e875ad376e3d75ddd1109ba800
+size 24704

temp/seg_map/seg_15.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4f3815fe49562a34bf0de78eacc0c0508300d35b8edbe27bbe3a58f10779747
+size 24704

temp/seg_map/seg_16.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c61b774d7ebf02ba7b6bfcda062cdf600f6d39f4d499b72d63bfa4ab566cb7f6
+size 24704

temp/seg_map/seg_17.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86f445d11dc69ea7c45cd409637f630abb56859cedd5204cd761f1a75aece7e9
+size 24704

temp/seg_map/seg_18.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19ef25b3f6cd2d76d62c96ccaec8f83c315eadbbba8c40cafb37c9a273d9cc8f
+size 24704