app.py

utils/utils.py

@@ -37,6 +37,28 @@ def load_cn_config(model_dir):
      config_path = os.path.join(os.getcwd(), file_name)
      shutil.copy(config_path, file_path)
 
+def load_tagger_model(model_dir):
+    model_id = 'SmilingWolf/wd-swinv2-tagger-v3'
+    files = [
+        'config.json', 'model.onnx', 'selected_tags.csv', 'sw_jax_cv_config.json'
+    ]
+
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    
+    for file in files:
+        file_path = os.path.join(model_dir, file)
+        if not os.path.exists(file_path):
+            url = f'https://huggingface.co/{model_id}/resolve/main/{file}'
+            response = requests.get(url, allow_redirects=True)
+            if response.status_code == 200:
+                with open(file_path, 'wb') as f:
+                    f.write(response.content)
+                print(f'Downloaded {file}')
+            else:
+                print(f'Failed to download {file}')
+        else:
+            print(f'{file} already exists.')
 
 
 def resize_image_aspect_ratio(image):