From 2793bcdf2f95ef4f81eeaa22b2e68beb0af803a4 Mon Sep 17 00:00:00 2001
From: anonstash <143359807+anonstash@users.noreply.github.com>
Date: Sun, 5 May 2024 10:24:29 -0400
Subject: [PATCH] add modelVersionId to model json (#279)

* add modelVersionId to model json

* fix by_hash referenced before assignment
---
 scripts/civitai_file_manage.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/scripts/civitai_file_manage.py b/scripts/civitai_file_manage.py
index 0162418..0927254 100644
--- a/scripts/civitai_file_manage.py
+++ b/scripts/civitai_file_manage.py
@@ -672,6 +672,7 @@ def find_and_save(api_response, sha256=None, file_name=None, json_file=None, no_
 
 def get_models(file_path, gen_hash=None):
     modelId = None
+    modelVersionId = None
     sha256 = None
     json_file = os.path.splitext(file_path)[0] + ".json"
     if os.path.exists(json_file):
@@ -681,12 +682,14 @@ def get_models(file_path, gen_hash=None):
                 
                 if 'modelId' in data:
                     modelId = data['modelId']
+                if 'modelVersionId' in data:
+                    modelVersionId = data['modelVersionId']
                 if 'sha256' in data and data['sha256']:
                     sha256 = data['sha256']
         except Exception as e:
             print(f"Failed to open {json_file}: {e}")
     
-    if not modelId or not sha256:
+    if not modelId or not modelVersionId or not sha256:
         if gen_hash:
             if not sha256:
                 sha256 = gen_sha256(file_path)
@@ -698,7 +701,7 @@ def get_models(file_path, gen_hash=None):
                 return None
     proxies, ssl = _api.get_proxies()
     try:
-        if not modelId:
+        if not modelId or not modelVersionId:
             response = requests.get(by_hash, timeout=(60,30), proxies=proxies, verify=ssl)
             if response.status_code == 200:
                 api_response = response.json()
@@ -707,6 +710,7 @@ def get_models(file_path, gen_hash=None):
                     return None
                 else:
                     modelId = api_response.get("modelId", "")
+                    modelVersionId = api_response.get("id", "")
             elif response.status_code == 503:
                 return "offline"
             elif response.status_code == 404:
@@ -719,6 +723,7 @@ def get_models(file_path, gen_hash=None):
                         data = json.load(f)
 
                     data['modelId'] = modelId
+                    data['modelVersionId'] = modelVersionId
                     data['sha256'] = sha256.upper()
                         
                     with open(json_file, 'w', encoding="utf-8") as f:
@@ -728,6 +733,7 @@ def get_models(file_path, gen_hash=None):
             else:
                 data = {
                     'modelId': modelId,
+                    'modelVersionId': modelVersionId,
                     'sha256': sha256.upper()
                     }
                 with open(json_file, 'w', encoding="utf-8") as f: