sd35 all-in-one safetensors support

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/model_quant.py

@@ -7,10 +7,10 @@ bnb = None
 quanto = None
 
 
-def create_bnb_config(kwargs):
+def create_bnb_config(kwargs = None):
     from modules import shared, devices
     if len(shared.opts.bnb_quantization) > 0:
-        if 'Model' in shared.opts.bnb_quantization and 'transformer' not in kwargs:
+        if 'Model' in shared.opts.bnb_quantization and 'transformer' not in (kwargs or {}):
             load_bnb()
             bnb_config = diffusers.BitsAndBytesConfig(
                 load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
@@ -19,8 +19,12 @@ def create_bnb_config(kwargs):
                 bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
                 bnb_4bit_compute_dtype=devices.dtype
             )
-            kwargs['quantization_config'] = bnb_config
             shared.log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
+            if kwargs is None:
+                return bnb_config
+            else:
+                kwargs['quantization_config'] = bnb_config
+                return kwargs
     return kwargs
 
 

modules/model_sd3.py

@@ -1,7 +1,7 @@
 import os
 import diffusers
 import transformers
-from modules import shared, devices, sd_models, sd_unet, model_te, model_quant
+from modules import shared, devices, sd_models, sd_unet, model_te, model_quant, model_tools
 
 
 def load_overrides(kwargs, cache_dir):
@@ -69,7 +69,7 @@ def load_quants(kwargs, repo_id, cache_dir):
 
 
 def load_missing(kwargs, fn, cache_dir):
-    keys = sd_models.get_safetensor_keys(fn)
+    keys = model_tools.get_safetensor_keys(fn)
     size = os.stat(fn).st_size // 1024 // 1024
     if size > 15000:
         repo_id = 'stabilityai/stable-diffusion-3.5-large'
@@ -129,7 +129,11 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
     if fn is not None and os.path.exists(fn) and os.path.isfile(fn):
         if fn.endswith('.safetensors'):
             loader = diffusers.StableDiffusion3Pipeline.from_single_file
-            kwargs = load_missing(kwargs, fn, cache_dir)
+            # required_modules = model_tools.get_modules(diffusers.StableDiffusion3Pipeline)
+            # have_modules = model_tools.get_safetensor_keys(fn)
+            # loaded_modules = model_tools.load_modules('stabilityai/stable-diffusion-3.5-medium', required_modules)
+            # kwargs = {**kwargs, **loaded_modules}
+            # kwargs = load_missing(kwargs, fn, cache_dir)
             repo_id = fn
         elif fn.endswith('.gguf'):
             kwargs = load_gguf(kwargs, fn)

modules/model_tools.py

@@ -0,0 +1,49 @@
+import inspect
+import diffusers
+import transformers
+import safetensors.torch
+from modules import shared, devices, model_quant
+
+
+def get_safetensor_keys(filename):
+    keys = []
+    try:
+        with safetensors.torch.safe_open(filename, framework="pt", device="cpu") as f:
+            keys = f.keys()
+    except Exception as e:
+        shared.log.error(f'Load dict: path="{filename}" {e}')
+    return keys
+
+
+def get_modules(model: callable):
+    signature = inspect.signature(model.__init__, follow_wrapped=True)
+    params = {param.name: param.annotation for param in signature.parameters.values() if param.annotation != inspect._empty and hasattr(param.annotation, 'from_pretrained')} # pylint: disable=protected-access
+    for name, cls in params.items():
+        shared.log.debug(f'Analyze: model={model} module={name} class={cls.__name__} loadable={getattr(cls, "from_pretrained", None)}')
+    return params
+
+
+def load_modules(repo_id: str, params: dict):
+    cache_dir = shared.opts.hfcache_dir
+    modules = {}
+    for name, cls in params.items():
+        subfolder = None
+        kwargs = {}
+        if cls == diffusers.AutoencoderKL:
+            subfolder = 'vae'
+        if cls == transformers.CLIPTextModel: # clip-vit-l
+            subfolder = 'text_encoder'
+        if cls == transformers.CLIPTextModelWithProjection: # clip-vit-g
+            subfolder = 'text_encoder_2'
+        if cls == transformers.T5EncoderModel: # t5-xxl
+            subfolder = 'text_encoder_3'
+            kwargs['quantization_config'] = model_quant.create_bnb_config()
+            kwargs['variant'] = 'fp16'
+        if cls == diffusers.SD3Transformer2DModel:
+            subfolder = 'transformer'
+            kwargs['quantization_config'] = model_quant.create_bnb_config()
+        if subfolder is None:
+            continue
+        shared.log.debug(f'Load: module={name} class={cls.__name__} repo={repo_id} location={subfolder}')
+        modules[name] = cls.from_pretrained(repo_id, subfolder=subfolder, cache_dir=cache_dir, torch_dtype=devices.dtype, **kwargs)
+    return modules

modules/sd_detect.py

@@ -84,6 +84,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
             if 'omnigen' in f.lower():
                 guess = 'OmniGen'
                 pipeline = 'custom'
+            if 'sd3' in f.lower():
+                guess = 'Stable Diffusion 3'
             if 'flux' in f.lower():
                 guess = 'FLUX'
                 if size > 11000 and size < 20000:

modules/sd_models.py

@@ -75,16 +75,6 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl
     return sd
 
 
-def get_safetensor_keys(filename):
-    keys = []
-    try:
-        with safetensors.torch.safe_open(filename, framework="pt", device="cpu") as f:
-            keys = f.keys()
-    except Exception as e:
-        shared.log.error(f'Load dict: path="{filename}" {e}')
-    return keys
-
-
 def get_state_dict_from_checkpoint(pl_sd):
     checkpoint_dict_replacements = {
         'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',