Add support for InstructPix2Pix (fixes #40; PR #60)

daam/trace.py

@@ -237,7 +237,12 @@ def _unravel_attn(self, x):
         with auto_autocast(dtype=torch.float32):
             for map_ in x:
                 map_ = map_.view(map_.size(0), h, w)
-                map_ = map_[map_.size(0) // 2:]  # Filter out unconditional
+                # For Instruct Pix2Pix, divide the map into three parts: text condition, image condition and unconditional,
+                # and only keep the text condition part, which is first of the three parts(as per diffusers implementation).
+                if map_.size(0) == 24:
+                    map_ = map_[:((map_.size(0) // 3)+1)]  # Filter out unconditional and image condition
+                else:
+                    map_ = map_[map_.size(0) // 2:] #  # Filter out unconditional
                 maps.append(map_)
 
         maps = torch.stack(maps, 0)  # shape: (tokens, heads, height, width)