Upload 3 files

dc.py

@@ -79,7 +79,7 @@ from env import (
     DOWNLOAD_VAE_LIST, DOWNLOAD_EMBEDS)
 
 from modutils import (to_list, list_uniq, list_sub, get_model_id_list, get_tupled_embed_list,
-                      get_tupled_model_list, get_lora_model_list, download_private_repo, download_things)
+                      get_tupled_model_list, get_lora_model_list, download_private_repo, download_things, download_link_model)
 
 # - **Download Models**
 download_model = ", ".join(DOWNLOAD_MODEL_LIST)
@@ -173,11 +173,15 @@ class GuiSD:
     def load_new_model(self, model_name, vae_model, task, controlnet_model, progress=gr.Progress(track_tqdm=True)):
 
         # download link model > model_name
+        if "http" in model_name: #
+            model_name, model_type = download_link_model(model_name, DIRECTORY_MODELS) #
+            is_link_model = True #
+        else: is_link_model = False #
 
         self.update_storage_models()
 
         vae_model = vae_model if vae_model != "None" else None
-        model_type = get_model_type(model_name)
+        model_type = get_model_type(model_name) if not is_link_model else model_type #
         dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
 
         if not os.path.exists(model_name):

env.py

@@ -51,6 +51,8 @@ LOAD_DIFFUSERS_FORMAT_MODEL = [
     'Spestly/OdysseyXL-4.0',
     'hanzogak/comradeshipXL-v14T14H',
     'hanzogak/comradeshipXL-v14VT',
+    'BlueDancer/Artisanica_XL',
+    'neta-art/neta-noob-1.0',
     'OnomaAIResearch/Illustrious-xl-early-release-v0',
     'Laxhar/noobai-XL-1.0',
     'Laxhar/noobai-XL-Vpred-1.0',

modutils.py

@@ -19,6 +19,9 @@ from unidecode import unidecode
 import copy
 from datetime import datetime, timezone, timedelta
 FILENAME_TIMEZONE = timezone(timedelta(hours=9)) # JST
+import torch
+from safetensors.torch import load_file
+import gc
 
 
 from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
@@ -1628,6 +1631,69 @@ def get_model_pipeline(repo_id: str):
         return default
 
 
+MODEL_TYPE_KEY = {
+    "model.diffusion_model.output_blocks.1.1.norm.bias": "SDXL",
+    "model.diffusion_model.input_blocks.11.0.out_layers.3.weight": "SD 1.5",
+    "double_blocks.0.img_attn.norm.key_norm.scale": "FLUX",
+    "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale": "FLUX",
+    "model.diffusion_model.joint_blocks.9.x_block.attn.ln_k.weight": "SD 3.5",
+}
+
+
+def safe_clean(path: str):
+    try:
+        if Path(path).exists():
+            if Path(path).is_dir(): shutil.rmtree(str(Path(path)))
+            else: Path(path).unlink()
+            print(f"Deleted: {path}")
+        else: print(f"File not found: {path}")
+    except Exception as e:
+        print(f"Failed to delete: {path} {e}")
+
+
+def read_safetensors_key(path: str):
+    try:
+        keys = []
+        state_dict = load_file(str(Path(path)))
+        for k in list(state_dict.keys()):
+            keys.append(k)
+            state_dict.pop(k)
+    except Exception as e:
+        print(e)
+    finally:
+        del state_dict
+        torch.cuda.empty_cache()
+        gc.collect()
+        return keys
+
+
+def get_model_type_from_key(path: str):
+    default = "SDXL"
+    try:
+        keys = read_safetensors_key(path)
+        for k, v in MODEL_TYPE_KEY.items():
+            if k in set(keys):
+                print(f"Model type is {v}.")
+                return v
+        print("Model type could not be identified.")
+    except Exception:
+        return default
+    return default
+
+
+def download_link_model(url: str, localdir: str):
+    try:
+        new_file = None
+        new_file = get_download_file(localdir, url, CIVITAI_API_KEY)
+        if not new_file or Path(new_file).suffix.lower() not in set([".safetensors", ".ckpt", ".bin", ".sft"]):
+            if Path(new_file).exists(): Path(new_file).unlink()
+            raise gr.Error(f"Safetensors file not found: {url}")
+        model_type = get_model_type_from_key(new_file)
+        return new_file, model_type
+    except Exception as e:
+        raise gr.Error(f"Failed to load single model file: {url} {e}")
+
+
 EXAMPLES_GUI = [
     [
         "1girl, souryuu asuka langley, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors, masterpiece, best quality, very aesthetic, absurdres",