Upload convert_url_to_diffusers_sdxl_gr.py

convert_url_to_diffusers_sdxl_gr.py

@@ -16,7 +16,7 @@ def is_repo_name(s):
     return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
 
 
-def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)):
+def download_thing(directory, url, civitai_api_key="", hf_token="", progress=gr.Progress(track_tqdm=True)):
     url = url.strip()
     if "drive.google.com" in url:
         original_dir = os.getcwd()
@@ -27,9 +27,11 @@ def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(trac
         url = url.replace("?download=true", "")
         if "/blob/" in url:
             url = url.replace("/blob/", "/resolve/")
-            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+        user_header = f'"Authorization: Bearer {hf_token}"'
+        if hf_token:
+            os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
         else:
-            os.system (f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+            os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
     elif "civitai.com" in url:
         if "?" in url:
             url = url.split("?")[0]
@@ -52,7 +54,7 @@ def get_local_model_list(dir_path):
     return model_list
 
 
-def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqdm=True)):
+def get_download_file(temp_dir, url, civitai_key, hf_token, progress=gr.Progress(track_tqdm=True)):
     if not "http" in url and is_repo_name(url) and not Path(url).exists():
         print(f"Use HF Repo: {url}")
         new_file = url
@@ -66,7 +68,7 @@ def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqd
         print(f"Start downloading: {url}")
         before = get_local_model_list(temp_dir)
         try:
-            download_thing(temp_dir, url.strip(), civitai_key)
+            download_thing(temp_dir, url.strip(), civitai_key, hf_token)
         except Exception:
             print(f"Download failed: {url}")
             return ""
@@ -215,11 +217,11 @@ def fuse_loras(pipe, lora_dict={}, temp_dir=".", civitai_key=""):
     pipe.unload_lora_weights()
 
 
-def convert_url_to_diffusers_sdxl(url, civitai_key="", is_upload_sf=False, half=True, vae=None,
+def convert_url_to_diffusers_sdxl(url, civitai_key="", hf_token="", is_upload_sf=False, half=True, vae=None,
                                   scheduler="Euler a", lora_dict={}, is_local=True, progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Start converting...")
     temp_dir = "."
-    new_file = get_download_file(temp_dir, url, civitai_key)
+    new_file = get_download_file(temp_dir, url, civitai_key, hf_token)
     if not new_file:
         print(f"Not found: {url}")
         return ""
@@ -273,15 +275,15 @@ def convert_url_to_diffusers_sdxl(url, civitai_key="", is_upload_sf=False, half=
     return new_repo_name
 
 
-def is_repo_exists(repo_id):
+def is_repo_exists(repo_id, hf_token):
     from huggingface_hub import HfApi
-    api = HfApi()
+    api = HfApi(token=hf_token)
     try:
         if api.repo_exists(repo_id=repo_id): return True
         else: return False
     except Exception as e:
         print(e)
-        print(f"Error: Failed to connect {repo_id}. ")
+        print(f"Error: Failed to connect {repo_id}.")
         return True # for safe
 
 
@@ -315,10 +317,9 @@ def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_ke
         print(f"Invalid user name: {hf_user}")
         progress(1, desc=f"Invalid user name: {hf_user}")
         return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
-    if hf_token and not os.environ.get("HF_TOKEN"): os.environ['HF_TOKEN'] = hf_token
     if not civitai_key and os.environ.get("CIVITAI_API_KEY"): civitai_key = os.environ.get("CIVITAI_API_KEY")
     lora_dict = {lora1: lora1s, lora2: lora2s, lora3: lora3s, lora4: lora4s, lora5: lora5s}
-    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, is_upload_sf, half, vae, scheduler, lora_dict, False)
+    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, hf_token, is_upload_sf, half, vae, scheduler, lora_dict, False)
     if not new_path: return ""
     new_repo_id = f"{hf_user}/{Path(new_path).stem}"
     if hf_repo != "": new_repo_id = f"{hf_user}/{hf_repo}"
@@ -326,7 +327,7 @@ def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_ke
         print(f"Invalid repo name: {new_repo_id}")
         progress(1, desc=f"Invalid repo name: {new_repo_id}")
         return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
-    if is_repo_exists(new_repo_id):
+    if is_repo_exists(new_repo_id, hf_token):
         print(f"Repo already exists: {new_repo_id}")
         progress(1, desc=f"Repo already exists: {new_repo_id}")
         return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")