[api] Allows dynamic add model zoo (#1397)

Change-Id: I3209f9d8a7b16915d6417013cba7e87b19e0e954
deepjavalibrary · Dec 1, 2021 · 11d9974 · 11d9974
1 parent efb4ff8
commit 11d9974
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java b/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
@@ -34,8 +34,7 @@ public abstract class ModelZoo {
     static {
         ServiceLoader<ZooProvider> providers = ServiceLoader.load(ZooProvider.class);
         for (ZooProvider provider : providers) {
-            ModelZoo zoo = provider.getModelZoo();
-            MODEL_ZOO_MAP.put(zoo.getGroupId(), zoo);
+            registerModelZoo(provider);
         }
     }
 
@@ -78,6 +77,16 @@ public ModelLoader getModelLoader(String name) {
      */
     public abstract Set<String> getSupportedEngines();
 
+    /**
+     * Refreshes model zoo.
+     *
+     * @param provider the {@code ZooProvider}
+     */
+    public static void registerModelZoo(ZooProvider provider) {
+        ModelZoo zoo = provider.getModelZoo();
+        MODEL_ZOO_MAP.putIfAbsent(zoo.getGroupId(), zoo);
+    }
+
     /**
      * Returns available model zoos.
      *