fred-dev/comfy_ui_ali

1import torch2from comfy.ldm.modules.attention import optimized_attention_for_device3import comfy.ops4 5class CLIPAttention(torch.nn.Module):6    def __init__(self, embed_dim, heads, dtype, device, operations):7        super().__init__()8 9        self.heads = heads10        self.q_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)11        self.k_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)12        self.v_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)13 14        self.out_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)15 16    def forward(self, x, mask=None, optimized_attention=None):17        q = self.q_proj(x)18        k = self.k_proj(x)19        v = self.v_proj(x)20 21        out = optimized_attention(q, k, v, self.heads, mask)22        return self.out_proj(out)23 24ACTIVATIONS = {"quick_gelu": lambda a: a * torch.sigmoid(1.702 * a),25               "gelu": torch.nn.functional.gelu,26               "gelu_pytorch_tanh": lambda a: torch.nn.functional.gelu(a, approximate="tanh"),27}28 29class CLIPMLP(torch.nn.Module):30    def __init__(self, embed_dim, intermediate_size, activation, dtype, device, operations):31        super().__init__()32        self.fc1 = operations.Linear(embed_dim, intermediate_size, bias=True, dtype=dtype, device=device)33        self.activation = ACTIVATIONS[activation]34        self.fc2 = operations.Linear(intermediate_size, embed_dim, bias=True, dtype=dtype, device=device)35 36    def forward(self, x):37        x = self.fc1(x)38        x = self.activation(x)39        x = self.fc2(x)40        return x41 42class CLIPLayer(torch.nn.Module):43    def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):44        super().__init__()45        self.layer_norm1 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)46        self.self_attn = CLIPAttention(embed_dim, heads, dtype, device, operations)47        self.layer_norm2 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)48        self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device, operations)49 50    def forward(self, x, mask=None, optimized_attention=None):51        x += self.self_attn(self.layer_norm1(x), mask, optimized_attention)52        x += self.mlp(self.layer_norm2(x))53        return x54 55 56class CLIPEncoder(torch.nn.Module):57    def __init__(self, num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):58        super().__init__()59        self.layers = torch.nn.ModuleList([CLIPLayer(embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations) for i in range(num_layers)])60 61    def forward(self, x, mask=None, intermediate_output=None):62        optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)63 64        if intermediate_output is not None:65            if intermediate_output < 0:66                intermediate_output = len(self.layers) + intermediate_output67 68        intermediate = None69        for i, l in enumerate(self.layers):70            x = l(x, mask, optimized_attention)71            if i == intermediate_output:72                intermediate = x.clone()73        return x, intermediate74 75class CLIPEmbeddings(torch.nn.Module):76    def __init__(self, embed_dim, vocab_size=49408, num_positions=77, dtype=None, device=None, operations=None):77        super().__init__()78        self.token_embedding = operations.Embedding(vocab_size, embed_dim, dtype=dtype, device=device)79        self.position_embedding = operations.Embedding(num_positions, embed_dim, dtype=dtype, device=device)80 81    def forward(self, input_tokens, dtype=torch.float32):82        return self.token_embedding(input_tokens, out_dtype=dtype) + comfy.ops.cast_to(self.position_embedding.weight, dtype=dtype, device=input_tokens.device)83 84 85class CLIPTextModel_(torch.nn.Module):86    def __init__(self, config_dict, dtype, device, operations):87        num_layers = config_dict["num_hidden_layers"]88        embed_dim = config_dict["hidden_size"]89        heads = config_dict["num_attention_heads"]90        intermediate_size = config_dict["intermediate_size"]91        intermediate_activation = config_dict["hidden_act"]92        num_positions = config_dict["max_position_embeddings"]93        self.eos_token_id = config_dict["eos_token_id"]94 95        super().__init__()96        self.embeddings = CLIPEmbeddings(embed_dim, num_positions=num_positions, dtype=dtype, device=device, operations=operations)97        self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)98        self.final_layer_norm = operations.LayerNorm(embed_dim, dtype=dtype, device=device)99 100    def forward(self, input_tokens=None, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=torch.float32):101        if embeds is not None:102            x = embeds + comfy.ops.cast_to(self.embeddings.position_embedding.weight, dtype=dtype, device=embeds.device)103        else:104            x = self.embeddings(input_tokens, dtype=dtype)105 106        mask = None107        if attention_mask is not None:108            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])109            mask = mask.masked_fill(mask.to(torch.bool), -torch.finfo(x.dtype).max)110 111        causal_mask = torch.full((x.shape[1], x.shape[1]), -torch.finfo(x.dtype).max, dtype=x.dtype, device=x.device).triu_(1)112 113        if mask is not None:114            mask += causal_mask115        else:116            mask = causal_mask117 118        x, i = self.encoder(x, mask=mask, intermediate_output=intermediate_output)119        x = self.final_layer_norm(x)120        if i is not None and final_layer_norm_intermediate:121            i = self.final_layer_norm(i)122 123        if num_tokens is not None:124            pooled_output = x[list(range(x.shape[0])), list(map(lambda a: a - 1, num_tokens))]125        else:126            pooled_output = x[torch.arange(x.shape[0], device=x.device), (torch.round(input_tokens).to(dtype=torch.int, device=x.device) == self.eos_token_id).int().argmax(dim=-1),]127        return x, i, pooled_output128 129class CLIPTextModel(torch.nn.Module):130    def __init__(self, config_dict, dtype, device, operations):131        super().__init__()132        self.num_layers = config_dict["num_hidden_layers"]133        self.text_model = CLIPTextModel_(config_dict, dtype, device, operations)134        embed_dim = config_dict["hidden_size"]135        self.text_projection = operations.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)136        self.dtype = dtype137 138    def get_input_embeddings(self):139        return self.text_model.embeddings.token_embedding140 141    def set_input_embeddings(self, embeddings):142        self.text_model.embeddings.token_embedding = embeddings143 144    def forward(self, *args, **kwargs):145        x = self.text_model(*args, **kwargs)146        out = self.text_projection(x[2])147        return (x[0], x[1], out, x[2])148 149 150class CLIPVisionEmbeddings(torch.nn.Module):151    def __init__(self, embed_dim, num_channels=3, patch_size=14, image_size=224, model_type="", dtype=None, device=None, operations=None):152        super().__init__()153 154        num_patches = (image_size // patch_size) ** 2155        if model_type == "siglip_vision_model":156            self.class_embedding = None157            patch_bias = True158        else:159            num_patches = num_patches + 1160            self.class_embedding = torch.nn.Parameter(torch.empty(embed_dim, dtype=dtype, device=device))161            patch_bias = False162 163        self.patch_embedding = operations.Conv2d(164            in_channels=num_channels,165            out_channels=embed_dim,166            kernel_size=patch_size,167            stride=patch_size,168            bias=patch_bias,169            dtype=dtype,170            device=device171        )172 173        self.position_embedding = operations.Embedding(num_patches, embed_dim, dtype=dtype, device=device)174 175    def forward(self, pixel_values):176        embeds = self.patch_embedding(pixel_values).flatten(2).transpose(1, 2)177        if self.class_embedding is not None:178            embeds = torch.cat([comfy.ops.cast_to_input(self.class_embedding, embeds).expand(pixel_values.shape[0], 1, -1), embeds], dim=1)179        return embeds + comfy.ops.cast_to_input(self.position_embedding.weight, embeds)180 181 182class CLIPVision(torch.nn.Module):183    def __init__(self, config_dict, dtype, device, operations):184        super().__init__()185        num_layers = config_dict["num_hidden_layers"]186        embed_dim = config_dict["hidden_size"]187        heads = config_dict["num_attention_heads"]188        intermediate_size = config_dict["intermediate_size"]189        intermediate_activation = config_dict["hidden_act"]190        model_type = config_dict["model_type"]191 192        self.embeddings = CLIPVisionEmbeddings(embed_dim, config_dict["num_channels"], config_dict["patch_size"], config_dict["image_size"], model_type=model_type, dtype=dtype, device=device, operations=operations)193        if model_type == "siglip_vision_model":194            self.pre_layrnorm = lambda a: a195            self.output_layernorm = True196        else:197            self.pre_layrnorm = operations.LayerNorm(embed_dim)198            self.output_layernorm = False199        self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)200        self.post_layernorm = operations.LayerNorm(embed_dim)201 202    def forward(self, pixel_values, attention_mask=None, intermediate_output=None):203        x = self.embeddings(pixel_values)204        x = self.pre_layrnorm(x)205        #TODO: attention_mask?206        x, i = self.encoder(x, mask=None, intermediate_output=intermediate_output)207        if self.output_layernorm:208            x = self.post_layernorm(x)209            pooled_output = x210        else:211            pooled_output = self.post_layernorm(x[:, 0, :])212        return x, i, pooled_output213 214class LlavaProjector(torch.nn.Module):215    def __init__(self, in_dim, out_dim, dtype, device, operations):216        super().__init__()217        self.linear_1 = operations.Linear(in_dim, out_dim, bias=True, device=device, dtype=dtype)218        self.linear_2 = operations.Linear(out_dim, out_dim, bias=True, device=device, dtype=dtype)219 220    def forward(self, x):221        return self.linear_2(torch.nn.functional.gelu(self.linear_1(x[:, 1:])))222 223class CLIPVisionModelProjection(torch.nn.Module):224    def __init__(self, config_dict, dtype, device, operations):225        super().__init__()226        self.vision_model = CLIPVision(config_dict, dtype, device, operations)227        if "projection_dim" in config_dict:228            self.visual_projection = operations.Linear(config_dict["hidden_size"], config_dict["projection_dim"], bias=False)229        else:230            self.visual_projection = lambda a: a231 232        if "llava3" == config_dict.get("projector_type", None):233            self.multi_modal_projector = LlavaProjector(config_dict["hidden_size"], 4096, dtype, device, operations)234        else:235            self.multi_modal_projector = None236 237    def forward(self, *args, **kwargs):238        x = self.vision_model(*args, **kwargs)239        out = self.visual_projection(x[2])240        projected = None241        if self.multi_modal_projector is not None:242            projected = self.multi_modal_projector(x[1])243 244        return (x[0], x[1], out, projected)245