From abfb197504c57d8e854dfd3e7d13537bc30d4a8e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 30 Oct 2024 19:30:52 -0400 Subject: [PATCH] sd35 all-in-one safetensors support Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 +++ modules/model_quant.py | 10 ++++++--- modules/model_sd3.py | 10 ++++++--- modules/model_tools.py | 49 ++++++++++++++++++++++++++++++++++++++++++ modules/sd_detect.py | 2 ++ modules/sd_models.py | 10 --------- 6 files changed, 68 insertions(+), 16 deletions(-) create mode 100644 modules/model_tools.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5deb1a928..71e8a667e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/modules/model_quant.py b/modules/model_quant.py index fab3a93a0..547a3d7ae 100644 --- a/modules/model_quant.py +++ b/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 diff --git a/modules/model_sd3.py b/modules/model_sd3.py index e369b197a..fc817dc4b 100644 --- a/modules/model_sd3.py +++ b/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) diff --git a/modules/model_tools.py b/modules/model_tools.py new file mode 100644 index 000000000..1212da244 --- /dev/null +++ b/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 diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 7144a7be7..148e788d5 100644 --- a/modules/sd_detect.py +++ b/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: diff --git a/modules/sd_models.py b/modules/sd_models.py index 3afc81e67..7743ce27d 100644 --- a/modules/sd_models.py +++ b/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.',