Merge pull request #164 from pytti-tools/test - MMC support

Author: dmarx

src/pytti/Notebook.py

@@ -18,6 +18,8 @@
 
 from pytti import Perceptor
 
+from pytti.Perceptor import CLIP_PERCEPTORS
+
 
 # https://stackoverflow.com/questions/15411967/how-can-i-check-if-code-is-executed-in-the-ipython-notebook
 def is_notebook():
@@ -248,11 +250,33 @@ def load_clip(params):
         if params.get(config_name):
             CLIP_MODEL_NAMES.append(clip_name)
 
-    if last_names != CLIP_MODEL_NAMES or Perceptor.CLIP_PERCEPTORS is None:
-        if CLIP_MODEL_NAMES == []:
+    if not params.get("use_mmc"):
+        if last_names != CLIP_MODEL_NAMES or Perceptor.CLIP_PERCEPTORS is None:
+            if CLIP_MODEL_NAMES == []:
+                Perceptor.free_clip()
+                raise RuntimeError("Please select at least one CLIP model")
             Perceptor.free_clip()
-            raise RuntimeError("Please select at least one CLIP model")
-        Perceptor.free_clip()
-        logger.debug("Loading CLIP...")
-        Perceptor.init_clip(CLIP_MODEL_NAMES)
-        logger.debug("CLIP loaded.")
+            logger.debug("Loading CLIP...")
+            Perceptor.init_clip(CLIP_MODEL_NAMES)
+            logger.debug("CLIP loaded.")
+    else:
+        logger.debug("attempting to use mmc to load perceptors")
+        import mmc
+        from mmc.registry import REGISTRY
+        import mmc.loaders  # force trigger model registrations
+        from mmc.mock.openai import MockOpenaiClip
+
+        CLIP_PERCEPTORS = (
+            []
+        )  # this will be fine because we'll use it to overwrite the module object
+        for item in params.mmc_models:
+            logger.debug(item)
+            model_loaders = REGISTRY.find(**item)
+            logger.debug(model_loaders)
+            for model_loader in model_loaders:
+                logger.debug(model_loader)
+                model = model_loader.load()
+                model = MockOpenaiClip(model)
+                CLIP_PERCEPTORS.append(model)
+        logger.debug(CLIP_PERCEPTORS)
+        Perceptor.CLIP_PERCEPTORS = CLIP_PERCEPTORS  # weird that htis works, but fine.

tests/test_mmc_loaders.py

@@ -0,0 +1,87 @@
+import pytest
+
+from hydra import initialize, compose
+from loguru import logger
+from pytti.workhorse import _main as render_frames
+from omegaconf import OmegaConf, open_dict
+
+
+CONFIG_BASE_PATH = "config"
+CONFIG_DEFAULTS = "default.yaml"
+
+
+cfg_yaml_openai = """# @package _global_
+scenes: a photograph of an apple
+use_mmc: true
+mmc_models:
+- architecture: clip
+  publisher: openai
+  id: RN50
+- architecture: clip
+  publisher: openai
+  id: ViT-B/32
+"""
+
+
+def test_mmc_openai_models():
+
+    with initialize(config_path=CONFIG_BASE_PATH):
+        cfg_base = compose(
+            config_name=CONFIG_DEFAULTS,
+            overrides=[f"conf=_empty"],
+        )
+        cfg_mmc = OmegaConf.create(cfg_yaml_openai)
+
+        with open_dict(cfg_base) as cfg:
+            cfg = OmegaConf.merge(cfg_base, cfg_mmc)
+        render_frames(cfg)
+
+
+cfg_yaml_mlf = """# @package _global_
+scenes: a photograph of an apple
+use_mmc: true
+mmc_models:
+- architecture: clip
+  publisher: mlfoundations
+  id: RN50--yfcc15m
+- architecture: clip
+  publisher: mlfoundations
+  id: ViT-B-32--laion400m_avg
+"""
+
+
+def test_mmc_mlf_models():
+
+    with initialize(config_path=CONFIG_BASE_PATH):
+        cfg_base = compose(
+            config_name=CONFIG_DEFAULTS,
+            overrides=[f"conf=_empty"],
+        )
+        cfg_mmc = OmegaConf.create(cfg_yaml_mlf)
+
+        with open_dict(cfg_base) as cfg:
+            cfg = OmegaConf.merge(cfg_base, cfg_mmc)
+        render_frames(cfg)
+
+
+cfg_yaml_all_cloob = """# @package _global_
+scenes: a photograph of an apple
+use_mmc: true
+mmc_models:
+- architecture: cloob
+  id: cloob_laion_400m_vit_b_16_32_epochs
+"""
+
+
+def test_mmc_all_cloob_models():
+
+    with initialize(config_path=CONFIG_BASE_PATH):
+        cfg_base = compose(
+            config_name=CONFIG_DEFAULTS,
+            overrides=[f"conf=_empty"],
+        )
+        cfg_mmc = OmegaConf.create(cfg_yaml_all_cloob)
+
+        with open_dict(cfg_base) as cfg:
+            cfg = OmegaConf.merge(cfg_base, cfg_mmc)
+        render_frames(cfg)