sd35 all-in-one safetensors support

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

CHANGELOG.md

@@ -9,6 +9,9 @@
   - *Alimama Inpainting, SoftEdge*
   - *note*: that just like with FLUX.1 or any large model, ControlNet are also large and can push your system over the limit  
     e.g. SD3 controlnets vary from 1GB to over 4GB in size  
+- SD3: support for all-in-one safetensors
+  - examples: [large](https://civitai.com/models/882666/sd35-large-google-flan?modelVersionId=1003031) [medium](https://civitai.com/models/900327)
+  - *note*: enable bnb on-the-fly quantization for even bigger gains  
 
 ## Update for 2024-10-29
 

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.',