Upload 3 files

app.py

@@ -65,7 +65,7 @@ from modutils import (list_uniq, download_private_repo, get_model_id_list, get_t
     get_lora_model_list, get_all_lora_tupled_list, update_loras, apply_lora_prompt, set_prompt_loras,
     get_my_lora, upload_file_lora, move_file_lora, search_civitai_lora, select_civitai_lora,
     update_civitai_selection, get_civitai_tag, CIVITAI_SORT, CIVITAI_PERIOD, CIVITAI_BASEMODEL,
-    set_textual_inversion_prompt, get_model_pipeline, change_interface_mode, get_t2i_model_info,
+    set_textual_inversion_prompt, get_model_pipeline, change_interface_mode, get_t2i_model_info, download_link_model,
     get_tupled_model_list, save_gallery_images, save_gallery_history, set_optimization, set_sampler_settings,
     set_quick_presets, process_style_prompt, optimization_list, save_images, download_things, valid_model_name,
     preset_styles, preset_quality, preset_sampler_setting, translate_to_en, EXAMPLES_GUI, RESOURCES)

env.py

@@ -21,6 +21,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',
     'Raelina/Rae-Diffusion-XL-V2',
     'Raelina/Raemu-XL-V4',

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",