model.py 15.4 KB
Newer Older
helloyongyang's avatar
helloyongyang committed
1
2
3
4
5
6
7
8
# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
PengGao's avatar
PengGao committed
9
from loguru import logger
helloyongyang's avatar
helloyongyang committed
10

helloyongyang's avatar
helloyongyang committed
11
12
# from lightx2v.attentions import attention
from lightx2v.common.ops.attn import TorchSDPAWeight
PengGao's avatar
PengGao committed
13
from lightx2v.models.input_encoders.hf.q_linear import Q8FQuantLinearFp8, Q8FQuantLinearInt8, TorchaoQuantLinearInt8, VllmQuantLinearFp8, VllmQuantLinearInt8
gushiqiao's avatar
gushiqiao committed
14
from lightx2v.utils.utils import load_weights
helloyongyang's avatar
helloyongyang committed
15
16

__all__ = [
Dongz's avatar
Dongz committed
17
18
19
    "XLMRobertaCLIP",
    "clip_xlm_roberta_vit_h_14",
    "CLIPModel",
helloyongyang's avatar
helloyongyang committed
20
21
22
23
24
25
26
27
28
29
]


def pos_interpolate(pos, seq_len):
    if pos.size(1) == seq_len:
        return pos
    else:
        src_grid = int(math.sqrt(pos.size(1)))
        tar_grid = int(math.sqrt(seq_len))
        n = pos.size(1) - src_grid * src_grid
Dongz's avatar
Dongz committed
30
31
32
33
34
35
36
        return torch.cat(
            [
                pos[:, :n],
                F.interpolate(pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(0, 3, 1, 2), size=(tar_grid, tar_grid), mode="bicubic", align_corners=False).flatten(2).transpose(1, 2),
            ],
            dim=1,
        )
helloyongyang's avatar
helloyongyang committed
37
38
39
40
41
42
43
44
45
46
47
48
49


class QuickGELU(nn.Module):
    def forward(self, x):
        return x * torch.sigmoid(1.702 * x)


class LayerNorm(nn.LayerNorm):
    def forward(self, x):
        return super().forward(x.float()).type_as(x)


class SelfAttention(nn.Module):
gushiqiao's avatar
gushiqiao committed
50
    def __init__(self, dim, num_heads, causal=False, attn_dropout=0.0, proj_dropout=0.0, quantized=False, quant_scheme=None, dtype=None):
helloyongyang's avatar
helloyongyang committed
51
52
53
54
55
56
57
58
59
60
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.causal = causal
        self.attn_dropout = attn_dropout
        self.proj_dropout = proj_dropout

        # layers
61
62
        if quantized:
            if quant_scheme == "int8":
gushiqiao's avatar
gushiqiao committed
63
                linear_cls = VllmQuantLinearInt8
64
            elif quant_scheme == "fp8":
gushiqiao's avatar
gushiqiao committed
65
66
67
                linear_cls = VllmQuantLinearFp8
            elif quant_scheme == "int8-torchao":
                linear_cls = TorchaoQuantLinearInt8
gushiqiao's avatar
gushiqiao committed
68
69
70
71
            elif quant_scheme == "int8-q8f":
                linear_cls = Q8FQuantLinearInt8
            elif quant_scheme == "fp8-q8f":
                linear_cls = Q8FQuantLinearFp8
72
73
74
        else:
            linear_cls = nn.Linear

gushiqiao's avatar
gushiqiao committed
75
76
        self.to_qkv = linear_cls(dim, dim * 3, dtype=dtype)
        self.proj = linear_cls(dim, dim, dtype=dtype)
helloyongyang's avatar
helloyongyang committed
77
78
79
80
81
82
83
84
85
86
87

    def forward(self, x):
        """
        x:   [B, L, C].
        """
        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim

        # compute query, key, value
        q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)

        # compute attention
helloyongyang's avatar
helloyongyang committed
88
        x = TorchSDPAWeight().apply(q=q, k=k, v=v)
helloyongyang's avatar
helloyongyang committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
        x = x.reshape(b, s, c)

        # output
        x = self.proj(x)
        x = F.dropout(x, self.proj_dropout, self.training)
        return x


class SwiGLU(nn.Module):
    def __init__(self, dim, mid_dim):
        super().__init__()
        self.dim = dim
        self.mid_dim = mid_dim
        # layers
        self.fc1 = nn.Linear(dim, mid_dim)
        self.fc2 = nn.Linear(dim, mid_dim)
        self.fc3 = nn.Linear(mid_dim, dim)

    def forward(self, x):
        x = F.silu(self.fc1(x)) * self.fc2(x)
        x = self.fc3(x)
        return x


class AttentionBlock(nn.Module):
gushiqiao's avatar
gushiqiao committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
    def __init__(
        self,
        dim,
        mlp_ratio,
        num_heads,
        post_norm=False,
        causal=False,
        activation="quick_gelu",
        attn_dropout=0.0,
        proj_dropout=0.0,
        norm_eps=1e-5,
        quantized=False,
        quant_scheme=None,
        dtype=torch.float16,
    ):
Dongz's avatar
Dongz committed
129
        assert activation in ["quick_gelu", "gelu", "swi_glu"]
helloyongyang's avatar
helloyongyang committed
130
131
132
133
134
135
136
137
138
        super().__init__()
        self.dim = dim
        self.mlp_ratio = mlp_ratio
        self.num_heads = num_heads
        self.post_norm = post_norm
        self.causal = causal
        self.norm_eps = norm_eps

        # layers
139
140
        if quantized:
            if quant_scheme == "int8":
gushiqiao's avatar
gushiqiao committed
141
                linear_cls = VllmQuantLinearInt8
142
            elif quant_scheme == "fp8":
gushiqiao's avatar
gushiqiao committed
143
144
145
                linear_cls = VllmQuantLinearFp8
            elif quant_scheme == "int8-torchao":
                linear_cls = TorchaoQuantLinearInt8
gushiqiao's avatar
gushiqiao committed
146
147
148
149
            elif quant_scheme == "int8-q8f":
                linear_cls = Q8FQuantLinearInt8
            elif quant_scheme == "fp8-q8f":
                linear_cls = Q8FQuantLinearFp8
150
151
152
        else:
            linear_cls = nn.Linear

gushiqiao's avatar
gushiqiao committed
153
154
155
        self.norm1 = LayerNorm(dim, eps=norm_eps, dtype=dtype)
        self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, proj_dropout, quantized, quant_scheme, dtype)
        self.norm2 = LayerNorm(dim, eps=norm_eps, dtype=dtype)
Dongz's avatar
Dongz committed
156
        if activation == "swi_glu":
gushiqiao's avatar
gushiqiao committed
157
            self.mlp = SwiGLU(dim, int(dim * mlp_ratio), dtype=dtype)
helloyongyang's avatar
helloyongyang committed
158
        else:
gushiqiao's avatar
gushiqiao committed
159
160
161
162
163
164
            self.mlp = nn.Sequential(
                linear_cls(dim, int(dim * mlp_ratio), dtype=dtype),
                QuickGELU() if activation == "quick_gelu" else nn.GELU(),
                linear_cls(int(dim * mlp_ratio), dim, dtype=dtype),
                nn.Dropout(proj_dropout),
            )
helloyongyang's avatar
helloyongyang committed
165
166
167
168
169
170
171
172
173
174
175
176

    def forward(self, x):
        if self.post_norm:
            x = x + self.norm1(self.attn(x))
            x = x + self.norm2(self.mlp(x))
        else:
            x = x + self.attn(self.norm1(x))
            x = x + self.mlp(self.norm2(x))
        return x


class AttentionPool(nn.Module):
gushiqiao's avatar
gushiqiao committed
177
    def __init__(self, dim, mlp_ratio, num_heads, activation="gelu", proj_dropout=0.0, norm_eps=1e-5, dtype=torch.float16):
helloyongyang's avatar
helloyongyang committed
178
179
180
181
182
183
184
185
186
187
188
189
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.mlp_ratio = mlp_ratio
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.proj_dropout = proj_dropout
        self.norm_eps = norm_eps

        # layers
        gain = 1.0 / math.sqrt(dim)
        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
gushiqiao's avatar
gushiqiao committed
190
191
192
193
194
195
196
        self.to_q = nn.Linear(dim, dim, dtype=dtype)
        self.to_kv = nn.Linear(dim, dim * 2, dtype=dtype)
        self.proj = nn.Linear(dim, dim, dtype=dtype)
        self.norm = LayerNorm(dim, eps=norm_eps, dtype=dtype)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio), dtype=dtype), QuickGELU() if activation == "quick_gelu" else nn.GELU(), nn.Linear(int(dim * mlp_ratio), dim, dtype=dtype), nn.Dropout(proj_dropout)
        )
helloyongyang's avatar
helloyongyang committed
197
198
199
200
201
202
203
204
205
206
207
208

    def forward(self, x):
        """
        x:  [B, L, C].
        """
        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim

        # compute query, key, value
        q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
        k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)

        # compute attention
Dongz's avatar
Dongz committed
209
        x = attention(q=q, k=k, v=v, attention_type="torch_sdpa")
helloyongyang's avatar
helloyongyang committed
210
211
212
213
214
215
216
217
218
219
220
221
        x = x.reshape(b, 1, c)

        # output
        x = self.proj(x)
        x = F.dropout(x, self.proj_dropout, self.training)

        # mlp
        x = x + self.mlp(self.norm(x))
        return x[:, 0]


class VisionTransformer(nn.Module):
Dongz's avatar
Dongz committed
222
223
    def __init__(
        self,
gushiqiao's avatar
gushiqiao committed
224
        dtype=torch.float16,
Dongz's avatar
Dongz committed
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
        image_size=224,
        patch_size=16,
        dim=768,
        mlp_ratio=4,
        out_dim=512,
        num_heads=12,
        num_layers=12,
        pool_type="token",
        pre_norm=True,
        post_norm=False,
        activation="quick_gelu",
        attn_dropout=0.0,
        proj_dropout=0.0,
        embedding_dropout=0.0,
        norm_eps=1e-5,
240
241
        quantized=False,
        quant_scheme=None,
Dongz's avatar
Dongz committed
242
    ):
helloyongyang's avatar
helloyongyang committed
243
        if image_size % patch_size != 0:
root's avatar
root committed
244
            logger.info("[WARNING] image_size is not divisible by patch_size", flush=True)
Dongz's avatar
Dongz committed
245
        assert pool_type in ("token", "token_fc", "attn_pool")
helloyongyang's avatar
helloyongyang committed
246
247
248
249
        out_dim = out_dim or dim
        super().__init__()
        self.image_size = image_size
        self.patch_size = patch_size
Dongz's avatar
Dongz committed
250
        self.num_patches = (image_size // patch_size) ** 2
helloyongyang's avatar
helloyongyang committed
251
252
253
254
255
256
257
258
259
260
261
        self.dim = dim
        self.mlp_ratio = mlp_ratio
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.pool_type = pool_type
        self.post_norm = post_norm
        self.norm_eps = norm_eps

        # embeddings
        gain = 1.0 / math.sqrt(dim)
gushiqiao's avatar
gushiqiao committed
262
        self.patch_embedding = nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size, bias=not pre_norm, dtype=dtype)
Dongz's avatar
Dongz committed
263
        if pool_type in ("token", "token_fc"):
gushiqiao's avatar
gushiqiao committed
264
265
            self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim, dtype=dtype))
        self.pos_embedding = nn.Parameter(gain * torch.randn(1, self.num_patches + (1 if pool_type in ("token", "token_fc") else 0), dim, dtype=dtype))
helloyongyang's avatar
helloyongyang committed
266
267
268
        self.dropout = nn.Dropout(embedding_dropout)

        # transformer
gushiqiao's avatar
gushiqiao committed
269
        self.pre_norm = LayerNorm(dim, eps=norm_eps, dtype=dtype) if pre_norm else None
270
        self.transformer = nn.Sequential(
gushiqiao's avatar
gushiqiao committed
271
            *[AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, activation, attn_dropout, proj_dropout, norm_eps, quantized, quant_scheme, dtype) for _ in range(num_layers)]
272
        )
gushiqiao's avatar
gushiqiao committed
273
        self.post_norm = LayerNorm(dim, eps=norm_eps, dtype=dtype)
helloyongyang's avatar
helloyongyang committed
274
275

        # head
Dongz's avatar
Dongz committed
276
        if pool_type == "token":
gushiqiao's avatar
gushiqiao committed
277
            self.head = nn.Parameter(gain * torch.randn(dim, out_dim, dtype=dtype))
Dongz's avatar
Dongz committed
278
        elif pool_type == "token_fc":
gushiqiao's avatar
gushiqiao committed
279
            self.head = nn.Linear(dim, out_dim, dtype=dtype)
Dongz's avatar
Dongz committed
280
        elif pool_type == "attn_pool":
gushiqiao's avatar
gushiqiao committed
281
            self.head = AttentionPool(dim, mlp_ratio, num_heads, activation, proj_dropout, norm_eps, dtype=dtype)
helloyongyang's avatar
helloyongyang committed
282
283
284
285
286
287

    def forward(self, x, interpolation=False, use_31_block=False):
        b = x.size(0)

        # embeddings
        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
Dongz's avatar
Dongz committed
288
        if self.pool_type in ("token", "token_fc"):
helloyongyang's avatar
helloyongyang committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
            x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)
        if interpolation:
            e = pos_interpolate(self.pos_embedding, x.size(1))
        else:
            e = self.pos_embedding
        x = self.dropout(x + e)
        if self.pre_norm is not None:
            x = self.pre_norm(x)

        # transformer
        if use_31_block:
            x = self.transformer[:-1](x)
            return x
        else:
            x = self.transformer(x)
            return x


class XLMRobertaCLIP(nn.Module):
Dongz's avatar
Dongz committed
308
309
    def __init__(
        self,
gushiqiao's avatar
gushiqiao committed
310
        dtype=torch.float16,
Dongz's avatar
Dongz committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
        embed_dim=1024,
        image_size=224,
        patch_size=14,
        vision_dim=1280,
        vision_mlp_ratio=4,
        vision_heads=16,
        vision_layers=32,
        vision_pool="token",
        vision_pre_norm=True,
        vision_post_norm=False,
        activation="gelu",
        vocab_size=250002,
        max_text_len=514,
        type_size=1,
        pad_id=1,
        attn_dropout=0.0,
        proj_dropout=0.0,
        embedding_dropout=0.0,
        norm_eps=1e-5,
330
331
        quantized=False,
        quant_scheme=None,
Dongz's avatar
Dongz committed
332
    ):
helloyongyang's avatar
helloyongyang committed
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
        super().__init__()
        self.embed_dim = embed_dim
        self.image_size = image_size
        self.patch_size = patch_size
        self.vision_dim = vision_dim
        self.vision_mlp_ratio = vision_mlp_ratio
        self.vision_heads = vision_heads
        self.vision_layers = vision_layers
        self.vision_pre_norm = vision_pre_norm
        self.vision_post_norm = vision_post_norm
        self.activation = activation
        self.vocab_size = vocab_size
        self.max_text_len = max_text_len
        self.type_size = type_size
        self.pad_id = pad_id
        self.norm_eps = norm_eps

        # models
        self.visual = VisionTransformer(
gushiqiao's avatar
gushiqiao committed
352
            dtype=dtype,
helloyongyang's avatar
helloyongyang committed
353
354
355
356
357
358
359
360
361
362
363
364
365
366
            image_size=image_size,
            patch_size=patch_size,
            dim=vision_dim,
            mlp_ratio=vision_mlp_ratio,
            out_dim=embed_dim,
            num_heads=vision_heads,
            num_layers=vision_layers,
            pool_type=vision_pool,
            pre_norm=vision_pre_norm,
            post_norm=vision_post_norm,
            activation=activation,
            attn_dropout=attn_dropout,
            proj_dropout=proj_dropout,
            embedding_dropout=embedding_dropout,
Dongz's avatar
Dongz committed
367
            norm_eps=norm_eps,
368
369
            quantized=quantized,
            quant_scheme=quant_scheme,
Dongz's avatar
Dongz committed
370
        )
helloyongyang's avatar
helloyongyang committed
371
372
373
        self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))


Dongz's avatar
Dongz committed
374
def _clip(pretrained=False, pretrained_name=None, model_cls=XLMRobertaCLIP, return_transforms=False, return_tokenizer=False, tokenizer_padding="eos", dtype=torch.float32, device="cpu", **kwargs):
helloyongyang's avatar
helloyongyang committed
375
376
    # init a model on device
    with torch.device(device):
gushiqiao's avatar
gushiqiao committed
377
        model = model_cls(dtype=dtype, **kwargs)
helloyongyang's avatar
helloyongyang committed
378

gushiqiao's avatar
gushiqiao committed
379
    model = model.to(device=device)
helloyongyang's avatar
helloyongyang committed
380

gushiqiao's avatar
gushiqiao committed
381
    output = (model,)
helloyongyang's avatar
helloyongyang committed
382
383
384
    # init transforms
    if return_transforms:
        # mean and std
Dongz's avatar
Dongz committed
385
        if "siglip" in pretrained_name.lower():
helloyongyang's avatar
helloyongyang committed
386
387
388
389
390
391
            mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
        else:
            mean = [0.48145466, 0.4578275, 0.40821073]
            std = [0.26862954, 0.26130258, 0.27577711]

        # transforms
Dongz's avatar
Dongz committed
392
        transforms = T.Compose([T.Resize((model.image_size, model.image_size), interpolation=T.InterpolationMode.BICUBIC), T.ToTensor(), T.Normalize(mean=mean, std=std)])
helloyongyang's avatar
helloyongyang committed
393
394
395
396
        output += (transforms,)
    return output[0] if len(output) == 1 else output


Dongz's avatar
Dongz committed
397
def clip_xlm_roberta_vit_h_14(pretrained=False, pretrained_name="open-clip-xlm-roberta-large-vit-huge-14", **kwargs):
helloyongyang's avatar
helloyongyang committed
398
399
400
401
402
403
404
405
    cfg = dict(
        embed_dim=1024,
        image_size=224,
        patch_size=14,
        vision_dim=1280,
        vision_mlp_ratio=4,
        vision_heads=16,
        vision_layers=32,
Dongz's avatar
Dongz committed
406
407
        vision_pool="token",
        activation="gelu",
helloyongyang's avatar
helloyongyang committed
408
409
410
411
412
413
        vocab_size=250002,
        max_text_len=514,
        type_size=1,
        pad_id=1,
        attn_dropout=0.0,
        proj_dropout=0.0,
Dongz's avatar
Dongz committed
414
415
        embedding_dropout=0.0,
    )
helloyongyang's avatar
helloyongyang committed
416
417
418
419
420
    cfg.update(**kwargs)
    return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)


class CLIPModel:
gushiqiao's avatar
gushiqiao committed
421
    def __init__(self, dtype, device, checkpoint_path, clip_quantized, clip_quantized_ckpt, quant_scheme, cpu_offload=False, use_31_block=True, seq_p_group=None):
helloyongyang's avatar
helloyongyang committed
422
423
        self.dtype = dtype
        self.device = device
424
        self.quantized = clip_quantized
gushiqiao's avatar
gushiqiao committed
425
426
        self.cpu_offload = cpu_offload
        self.use_31_block = use_31_block
427
428
        self.seq_p_group = seq_p_group

429
430
431
432
433
        if self.quantized:
            self.checkpoint_path = clip_quantized_ckpt
        else:
            self.checkpoint_path = checkpoint_path

helloyongyang's avatar
helloyongyang committed
434
        # init model
435
436
437
        self.model, self.transforms = clip_xlm_roberta_vit_h_14(
            pretrained=False, return_transforms=True, return_tokenizer=False, dtype=dtype, device=device, quantized=self.quantized, quant_scheme=quant_scheme
        )
helloyongyang's avatar
helloyongyang committed
438
        self.model = self.model.eval().requires_grad_(False)
gushiqiao's avatar
gushiqiao committed
439
        weight_dict = load_weights(self.checkpoint_path, cpu_offload=cpu_offload, remove_key="textual")
440
        self.model.load_state_dict(weight_dict)
helloyongyang's avatar
helloyongyang committed
441

gushiqiao's avatar
gushiqiao committed
442
443
    def visual(self, videos):
        if self.cpu_offload:
gushiqiao's avatar
gushiqiao committed
444
            self.to_cuda()
helloyongyang's avatar
helloyongyang committed
445
446
        # preprocess
        size = (self.model.image_size,) * 2
wangshankun's avatar
wangshankun committed
447
        videos = torch.cat([F.interpolate(u, size=size, mode="bicubic", align_corners=False) for u in videos])
helloyongyang's avatar
helloyongyang committed
448
449
        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
        # forward
Dongz's avatar
Dongz committed
450
        with torch.amp.autocast("cuda", dtype=self.dtype):
gushiqiao's avatar
gushiqiao committed
451
            out = self.model.visual(videos, use_31_block=self.use_31_block)
gushiqiao's avatar
gushiqiao committed
452

gushiqiao's avatar
gushiqiao committed
453
        if self.cpu_offload:
gushiqiao's avatar
gushiqiao committed
454
455
456
457
458
459
460
461
            self.to_cpu()
        return out

    def to_cuda(self):
        self.model = self.model.cuda()

    def to_cpu(self):
        self.model = self.model.cpu()