pyiqa.archs.afineclip_model

Module Contents

pyiqa.archs.afineclip_model.available_models() → List[str][source]

Returns the names of available CLIP models

pyiqa.archs.afineclip_model.load(name: str, device: Union[str, torch.device]='cuda' if torch.cuda.is_available() else 'cpu', jit: bool = False, download_root: str = None)[source]

Load a CLIP model

Parameters:

  • name (str) – A model name listed by clip.available_models(), or the path to a model checkpoint containing the state_dict

  • device (Union[str, torch.device]) – The device to put the loaded model

  • jit (bool) – Whether to load the optimized JIT model or more hackable non-JIT model (default).

  • download_root (str) – path to download the model files; by default, it uses “~/.cache/clip”

Returns:

  • model (torch.nn.Module) – The CLIP model

  • preprocess (Callable[[PIL.Image], torch.Tensor]) – A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input