From 5ab99d589a1f816f49410d8a721cb7fa79135d80 Mon Sep 17 00:00:00 2001
From: BlafKing <tombon59@gmail.com>
Date: Tue, 21 May 2024 16:00:11 +0200
Subject: [PATCH] Small hotfixes

---
 install.py                     | 11 +++++------
 requirements.txt               |  5 +++++
 scripts/civitai_api.py         | 19 +++++++------------
 scripts/civitai_file_manage.py | 19 +++++++++----------
 4 files changed, 26 insertions(+), 28 deletions(-)
 create mode 100644 requirements.txt

diff --git a/install.py b/install.py
index ffb0b0e..7132eaf 100644
--- a/install.py
+++ b/install.py
@@ -12,9 +12,8 @@ def install_req(check_name, install_name=None):
     if not launch.is_installed(f"{check_name}"):
         launch.run_pip(f"install {install_name}", "requirements for CivitAI Browser")
 
-install_req("send2trash")
-install_req("zip_unicode", "ZipUnicode")
-install_req("bs4", "beautifulsoup4")
-install_req("fake_useragent")
-install_req("packaging")
-install_req("pysocks")
\ No newline at end of file
+install_req("send2trash==1.8.2")
+install_req("zip_unicode", "ZipUnicode==1.1.1")
+install_req("bs4", "beautifulsoup4==4.12.3")
+install_req("packaging==23.2")
+install_req("pysocks==1.7.1")
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..1bbc4bb
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+send2trash==1.8.2
+ZipUnicode==1.1.1
+beautifulsoup4==4.12.3
+packaging==23.2
+pysocks==1.7.1
\ No newline at end of file
diff --git a/scripts/civitai_api.py b/scripts/civitai_api.py
index bc74ab0..a8b44be 100644
--- a/scripts/civitai_api.py
+++ b/scripts/civitai_api.py
@@ -19,10 +19,6 @@ from html import escape
 from scripts.civitai_global import print, debug_print
 import scripts.civitai_global as gl
 import scripts.civitai_download as _download
-try:
-    from fake_useragent import UserAgent
-except ImportError:
-    print("Python module 'fake_useragent' has not been imported correctly, please try to restart or install it manually.")
 
 gl.init()
 
@@ -66,6 +62,7 @@ def contenttype_folder(content_type, desc=None, fromCheck=False, custom_folder=N
             folder = os.path.join(extensions_dir, "stable-diffusion-webui-aesthetic-gradients", "aesthetic_embeddings")
         else:
             folder = os.path.join(custom_folder, "aesthetic_embeddings")
+            
     elif content_type == "LORA":
         if cmd_opts.lora_dir and not custom_folder:
             folder = cmd_opts.lora_dir
@@ -92,8 +89,11 @@ def contenttype_folder(content_type, desc=None, fromCheck=False, custom_folder=N
         else:
             folder = os.path.join(main_models, "VAE")
             
-    elif content_type == "Controlnet":  
-        folder = os.path.join(main_models, "ControlNet")
+    elif content_type == "Controlnet":
+        if cmd_opts.controlnet_dir and not custom_folder:
+            folder = cmd_opts.controlnet_dir
+        else:
+            folder = os.path.join(main_models, "ControlNet")
             
     elif content_type == "Poses":
         folder = os.path.join(main_models, "Poses")
@@ -1165,16 +1165,11 @@ def get_proxies():
     return proxies, ssl
 
 def get_headers(referer=None, no_api=None):
-    
     api_key = getattr(opts, "custom_api_key", "")
-    try:
-        user_agent = UserAgent().chrome
-    except ImportError:
-        user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36"
     headers = {
         "Connection": "keep-alive",
         "Sec-Ch-Ua-Platform": "Windows",
-        "User-Agent": user_agent,
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36",
         "Content-Type": "application/json"
     }
     if referer:
diff --git a/scripts/civitai_file_manage.py b/scripts/civitai_file_manage.py
index 0927254..be86cee 100644
--- a/scripts/civitai_file_manage.py
+++ b/scripts/civitai_file_manage.py
@@ -690,15 +690,14 @@ def get_models(file_path, gen_hash=None):
             print(f"Failed to open {json_file}: {e}")
     
     if not modelId or not modelVersionId or not sha256:
-        if gen_hash:
-            if not sha256:
-                sha256 = gen_sha256(file_path)
+        if not sha256 and gen_hash:
+            sha256 = gen_sha256(file_path)
+        
+        if sha256:
             by_hash = f"https://civitai.com/api/v1/model-versions/by-hash/{sha256}"
         else:
-            if modelId:
-                return modelId
-            else:
-                return None
+            return modelId if modelId else None
+
     proxies, ssl = _api.get_proxies()
     try:
         if not modelId or not modelVersionId:
@@ -714,9 +713,9 @@ def get_models(file_path, gen_hash=None):
             elif response.status_code == 503:
                 return "offline"
             elif response.status_code == 404:
-                api_response = response.json()
-                modelId = api_response.get("error", "")
-            
+                modelId = "Model not found"
+                modelVersionId = "Model not found"
+                
             if os.path.exists(json_file):
                 try:
                     with open(json_file, 'r', encoding="utf-8") as f: