Function _clip

diffsynth/models/wan_video_image_encoder.py:738–819

(pretrained=False,
          pretrained_name=None,
          model_cls=CLIP,
          return_transforms=False,
          return_tokenizer=False,
          tokenizer_padding='eos',
          dtype=torch.float32,
          device='cpu',
          **kwargs)

def _clip(pretrained=False,
          pretrained_name=None,
          model_cls=CLIP,
          return_transforms=False,
          return_tokenizer=False,
          tokenizer_padding='eos',
          dtype=torch.float32,
          device='cpu',
          **kwargs):
    # init model
    if pretrained and pretrained_name:
        from sora import BUCKET, DOWNLOAD_TO_CACHE

        # init a meta model
        with torch.device('meta'):
            model = model_cls(**kwargs)

        # checkpoint path
        checkpoint = f'models/clip/{pretrained_name}'
        if dtype in (torch.float16, torch.bfloat16):
            suffix = '-' + {
                torch.float16: 'fp16',
                torch.bfloat16: 'bf16'
            }[dtype]
            if object_exists(BUCKET, f'{checkpoint}{suffix}.pth'):
                checkpoint = f'{checkpoint}{suffix}'
        checkpoint += '.pth'

        # load
        model.load_state_dict(
            torch.load(DOWNLOAD_TO_CACHE(checkpoint), map_location=device),
            assign=True,
            strict=False)
    else:
        # init a model on device
        with torch.device(device):
            model = model_cls(**kwargs)

    # set device
    output = (model,)

    # init transforms
    if return_transforms:
        # mean and std
        if 'siglip' in pretrained_name.lower():
            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
        transforms = T.Compose([
            T.Resize((model.image_size, model.image_size),
                     interpolation=T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std)
        ])
        output += (transforms,)