Move src/clip -> src/open_clip, create setup.py + MANIFEST.in and tes…

…t package install as `open_clip`

MANIFEST.in

@@ -0,0 +1,3 @@
+include src/open_clip/bpe_simple_vocab_16e6.txt.gz
+include src/open_clip/model_configs/*.json
+

setup.py

@@ -0,0 +1,56 @@
+""" Setup
+"""
+from setuptools import setup, find_packages
+from codecs import open
+from os import path
+
+here = path.abspath(path.dirname(__file__))
+
+# Get the long description from the README file
+with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+exec(open('src/open_clip/version.py').read())
+setup(
+    name='open_clip',
+    version=__version__,
+    description='OpenCLIP',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/mlfoundations/open_clip',
+    author='',
+    author_email='',
+    classifiers=[
+        # How mature is this project? Common values are
+        #   3 - Alpha
+        #   4 - Beta
+        #   5 - Production/Stable
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+
+    # Note that this is a string of words separated by whitespace, not a list.
+    keywords='CLIP pretrained',
+    package_dir={'': 'src'},
+    packages=find_packages(where='src', exclude=['training']),
+    include_package_data=True,
+    install_requires=[
+        'torch >= 1.9',
+        'torchvision',
+        'webdataset >= 0.2.5',
+        'ftfy',
+        'regex',
+    ],
+    python_requires='>=3.7',
+)

src/clip/__init__.py → src/open_clip/__init__.py

@@ -1,7 +1,7 @@
 from .factory import create_model_and_transforms
 from .loss import ClipLoss
 from .model import CLIP, CLIPTextCfg, CLIPVisionCfg, convert_weights_to_fp16
-from .openai import load_openai
+from .openai import load_openai_model, list_openai_models
 from .pretrained import list_pretrained, list_pretrained_tag_models, list_pretrained_model_tags,\
     get_pretrained_url, download_pretrained
 from .tokenizer import SimpleTokenizer, tokenize

src/clip/factory.py → src/open_clip/factory.py

@@ -6,7 +6,7 @@
 import torch
 
 from .model import CLIP, convert_weights_to_fp16
-from .openai import load_openai
+from .openai import load_openai_model
 from .pretrained import get_pretrained_url, download_pretrained
 from .transform import image_transform
 
@@ -25,14 +25,14 @@ def load_state_dict(checkpoint_path: str, map_location='cpu'):
 def create_model_and_transforms(
         model_name: str,
         pretrained: str,
-        precision: str,
-        device: torch.device,
+        precision: str = 'fp32',
+        device: torch.device = torch.device('cpu'),
         force_quick_gelu: bool = False,
 ):
     pretrained = pretrained.lower()
     if pretrained == 'openai':
         logging.info(f'Loading pretrained {model_name} from OpenAI.')
-        model, preprocess_train, preprocess_val = load_openai(model_name, device=device, jit=False)
+        model, preprocess_train, preprocess_val = load_openai_model(model_name, device=device, jit=False)
         # See https://discuss.pytorch.org/t/valueerror-attemting-to-unscale-fp16-gradients/81372
         if precision == "amp" or precision == "fp32":
             model = model.float()
@@ -67,6 +67,7 @@ def create_model_and_transforms(
 
         model.to(device=device)
         if precision == "fp16":
+            assert device.type != 'cpu'
             convert_weights_to_fp16(model)
 
     return model, preprocess_train, preprocess_val

src/clip/openai.py → src/open_clip/openai.py

@@ -9,15 +9,15 @@
 from .pretrained import get_pretrained_url, list_pretrained_tag_models, download_pretrained
 from .transform import image_transform
 
-__all__ = ["available_openai_models", "load_openai"]
+__all__ = ["list_openai_models", "load_openai_model"]
 
 
-def available_openai_models() -> List[str]:
+def list_openai_models() -> List[str]:
     """Returns the names of available CLIP models"""
     return list_pretrained_tag_models('openai')
 
 
-def load_openai(
+def load_openai_model(
         name: str,
         device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
         jit=True,
@@ -45,7 +45,7 @@ def load_openai(
     elif os.path.isfile(name):
         model_path = name
     else:
-        raise RuntimeError(f"Model {name} not found; available models = {available_openai_models()}")
+        raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
 
     try:
         # loading JIT archive

src/clip/pretrained.py → src/open_clip/pretrained.py

@@ -49,8 +49,11 @@
 }
 
 
-def list_pretrained(as_tuple: bool = False):
-    return [(k, t) if as_tuple else ':'.join([k, t]) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
+def list_pretrained(as_str: bool = False):
+    """ returns list of pretrained models
+    Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
+    """
+    return [':'.join([k, t]) if as_str else (k, t) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
 
 
 def list_pretrained_tag_models(tag: str):

src/clip/transform.py → src/open_clip/transform.py

@@ -27,4 +27,4 @@ def image_transform(
             _convert_to_rgb,
             ToTensor(),
             normalize,
-        ])
+        ])

src/open_clip/version.py

@@ -0,0 +1 @@
+__version__ = '0.2.0'

src/training/data.py

@@ -22,7 +22,7 @@
 except ImportError:
     hvd = None
 
-from clip import tokenize
+from open_clip import tokenize
 
 
 class CsvDataset(Dataset):

src/training/main.py

@@ -10,20 +10,19 @@
 from torch import optim
 from torch.cuda.amp import GradScaler
 from torch.utils.tensorboard import SummaryWriter
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
 
-from clip.factory import create_model_and_transforms
+from open_clip import create_model_and_transforms
 from training.data import get_data
 from training.distributed import is_master, init_distributed_device, world_info_from_env
 from training.logger import setup_logging
 from training.params import parse_args
 from training.scheduler import cosine_lr
 from training.train import train_one_epoch, evaluate
 
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
 
 def random_seed(seed=42, rank=0):
     torch.manual_seed(seed + rank)

src/training/train.py

@@ -10,7 +10,7 @@
 import torch.nn.functional as F
 import wandb
 
-from clip.loss import ClipLoss
+from open_clip import ClipLoss
 from .distributed import is_master
 from .zero_shot import zero_shot_eval
 

src/training/zero_shot.py

@@ -5,7 +5,7 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from clip import tokenize
+from open_clip import tokenize
 from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template