Method new_forward

src/model.py:231–266

(self, x, context=None, training=False, mask=None)

        change_checkpoint(self.model.diffusion_model)

        def new_forward(self, x, context=None, training=False, mask=None):
            h = self.heads
            crossattn = False
            if context is not None:
                crossattn = True
            q = self.to_q(x)
            context = default(context, x)
            k = self.to_k(context)
            v = self.to_v(context)
            if crossattn:
                modifier = torch.ones_like(k)
                modifier[:, :1, :] = modifier[:, :1, :]*0.
                k = modifier*k + (1-modifier)*k.detach()
                v = modifier*v + (1-modifier)*v.detach()

            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
            sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
            attn = sim.softmax(dim=-1)

            if crossattn and (attention_store_background_mask or attention_store_step):
                if attn16:
                    if attn.shape[1] == 256:
                        controller(attn) 
                else:
                    if attn.shape[1] in store_list:
                        controller(attn) 
            elif crossattn and check_token_attn:
                if attn.shape[1] == 64:
                    place = 'mid'
                else:
                    place = 'down'
                controller(attn, crossattn, place)
            
            out = einsum('b i j, b j d -> b i d', attn, v)
            out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
            return self.to_out(out)
        cross_att_count = 0
        def change_forward(model, count):
            for layer in model.children():