Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

on-the-fly quant for sd35, flux and sana

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-12-20 17:28:32 -05:00
parent 1649e4e3c4
commit 58ad18ee58
8 changed files with 113 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_flux.py
+27 -6
View File
@@ -110,6 +110,7 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
return transformer, text_encoder_2
"""
def quant_flux_bnb(checkpoint_info, transformer, text_encoder_2):
repo_id = sd_models.path_to_repo(checkpoint_info.name)
cache_dir=shared.opts.diffusers_dir
@@ -139,6 +140,24 @@ def quant_flux_bnb(checkpoint_info, transformer, text_encoder_2):
from modules import errors
errors.display(e, 'FLUX:')
return transformer, text_encoder_2
"""
def load_quants(kwargs, repo_id, cache_dir):
if len(shared.opts.bnb_quantization) > 0:
quant_args = {}
quant_args = model_quant.create_bnb_config(quant_args)
quant_args = model_quant.create_ao_config(quant_args)
if not quant_args:
return
model_quant.load_bnb(f'Load model: type=FLUX quant={quant_args}')
if 'Model' in shared.opts.bnb_quantization and 'transformer' not in kwargs:
kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
if 'Text Encoder' in shared.opts.bnb_quantization and 'text_encoder_3' not in kwargs:
kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
return kwargs
def load_flux_gguf(file_path):
@@ -148,9 +167,8 @@ def load_flux_gguf(file_path):
from diffusers.loaders.single_file_utils import convert_flux_transformer_checkpoint_to_diffusers
from modules import ggml, sd_hijack_accelerate
with init_empty_weights():
from diffusers import FluxTransformer2DModel
config = FluxTransformer2DModel.load_config(os.path.join('configs', 'flux'), subfolder="transformer")
transformer = FluxTransformer2DModel.from_config(config).to(devices.dtype)
config = diffusers.FluxTransformer2DModel.load_config(os.path.join('configs', 'flux'), subfolder="transformer")
transformer = diffusers.FluxTransformer2DModel.from_config(config).to(devices.dtype)
expected_state_dict_keys = list(transformer.state_dict().keys())
state_dict, stats = ggml.load_gguf_state_dict(file_path, devices.dtype)
state_dict = convert_flux_transformer_checkpoint_to_diffusers(state_dict)
@@ -295,7 +313,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
# initialize pipeline with pre-loaded components
kwargs = {}
# transformer, text_encoder_2 = quant_flux_bnb(checkpoint_info, transformer, text_encoder_2)
if transformer is not None:
kwargs['transformer'] = transformer
sd_unet.loaded_unet = shared.opts.sd_unet
@@ -324,10 +341,14 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
kwargs[c] = kwargs[c].to(dtype=devices.dtype)
allow_quant = 'gguf' not in (sd_unet.loaded_unet or '')
fn = checkpoint_info.path
if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
# transformer, text_encoder_2 = quant_flux_bnb(checkpoint_info, transformer, text_encoder_2)
kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir)
kwargs = model_quant.create_bnb_config(kwargs, allow_quant)
kwargs = model_quant.create_ao_config(kwargs, allow_quant)
if checkpoint_info.path.endswith('.safetensors') and os.path.isfile(checkpoint_info.path):
pipe = diffusers.FluxPipeline.from_single_file(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
if fn.endswith('.safetensors') and os.path.isfile(fn):
pipe = diffusers.FluxPipeline.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
else:
pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
modules/model_sana.py
+43 -24
View File
@@ -1,54 +1,73 @@
import os
import time
import torch
import diffusers
import transformers
from modules import shared, sd_models, devices, modelloader, model_quant
"""
Efficient-Large-Model/Sana_1600M_1024px_MultiLing_diffusers
Efficient-Large-Model/Sana_1600M_1024px_diffusers
Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers
Efficient-Large-Model/Sana_1600M_512px_MultiLing_diffusers
Efficient-Large-Model/Sana_1600M_512px_diffusers
Efficient-Large-Model/Sana_600M_1024px_diffusers
Efficient-Large-Model/Sana_600M_512px_diffusers
"""
def load_quants(kwargs, repo_id, cache_dir):
if len(shared.opts.bnb_quantization) > 0:
quant_args = {}
quant_args = model_quant.create_bnb_config(quant_args)
quant_args = model_quant.create_ao_config(quant_args)
load_args = kwargs.copy()
if not quant_args:
return
model_quant.load_bnb(f'Load model: type=SD3 quant={quant_args} args={load_args}')
if 'Model' in shared.opts.bnb_quantization and 'transformer' not in kwargs:
kwargs['transformer'] = diffusers.models.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
if 'Text Encoder' in shared.opts.bnb_quantization and 'text_encoder_3' not in kwargs:
kwargs['text_encoder_3'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args)
shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
return kwargs
def load_sana(checkpoint_info, kwargs={}):
from modules import shared, sd_models, devices, modelloader, model_quant
modelloader.hf_login()
repo_id = checkpoint_info if isinstance(checkpoint_info, str) else checkpoint_info.path
repo_id = sd_models.path_to_repo(repo_id)
fn = checkpoint_info if isinstance(checkpoint_info, str) else checkpoint_info.path
repo_id = sd_models.path_to_repo(fn)
kwargs.pop('load_connected_pipeline', None)
kwargs.pop('safety_checker', None)
kwargs.pop('requires_safety_checker', None)
kwargs.pop('torch_dtype', None)
if not repo_id.endswith('_diffusers'):
repo_id = f'{repo_id}_diffusers'
if devices.dtype == torch.bfloat16 and 'BF16' not in repo_id:
repo_id = repo_id.replace('_diffusers', '_BF16_diffusers')
if 'Sana_1600M' in repo_id:
if devices.dtype == torch.bfloat16:
repo_id = 'Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers'
if devices.dtype == torch.bfloat16 or 'BF16' in repo_id:
if 'BF16' not in repo_id:
repo_id = repo_id.replace('_diffusers', '_BF16_diffusers')
kwargs['variant'] = 'bf16'
kwargs['torch_dtype'] = devices.dtype
else:
repo_id = 'Efficient-Large-Model/Sana_1600M_1024px_diffusers'
kwargs['variant'] = 'fp16'
if 'Sana_600M' in repo_id:
repo_id = 'Efficient-Large-Model/Sana_600M_1024px_diffusers'
kwargs['variant'] = 'fp16'
kwargs = model_quant.create_bnb_config(kwargs)
kwargs = model_quant.create_ao_config(kwargs)
shared.log.debug(f'Load model: type=Sana repo="{repo_id}" args={kwargs}')
if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir)
# kwargs = model_quant.create_bnb_config(kwargs)
# kwargs = model_quant.create_ao_config(kwargs)
shared.log.debug(f'Load model: type=Sana repo="{repo_id}" args={list(kwargs)}')
t0 = time.time()
pipe = diffusers.SanaPipeline.from_pretrained(repo_id, cache_dir = shared.opts.diffusers_dir, **kwargs)
pipe = diffusers.SanaPipeline.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs)
if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
pipe.text_encoder = pipe.text_encoder.to(dtype=devices.dtype)
if 'transformer' not in kwargs:
pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
if 'text_encoder' not in kwargs:
pipe.text_encoder = pipe.text_encoder.to(dtype=devices.dtype)
pipe.vae = pipe.vae.to(dtype=devices.dtype)
if devices.dtype == torch.float16:
pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float32) # gemma2 does not support fp16
if 'transformer' not in kwargs:
pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
if 'text_encoder' not in kwargs:
pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float32) # gemma2 does not support fp16
pipe.vae = pipe.vae.to(dtype=torch.float32) # dc-ae often overflows in fp16
if shared.opts.diffusers_eval:
pipe.text_encoder.eval()
modules/model_sd3.py
+9 -11
View File
@@ -51,19 +51,17 @@ def load_overrides(kwargs, cache_dir):
def load_quants(kwargs, repo_id, cache_dir):
if len(shared.opts.bnb_quantization) > 0:
model_quant.load_bnb('Load model: type=SD3')
bnb_config = diffusers.BitsAndBytesConfig(
load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
bnb_4bit_compute_dtype=devices.dtype
)
quant_args = {}
quant_args = model_quant.create_bnb_config(quant_args)
quant_args = model_quant.create_ao_config(quant_args)
if not quant_args:
return
model_quant.load_bnb(f'Load model: type=SD3 quant={quant_args}')
if 'Model' in shared.opts.bnb_quantization and 'transformer' not in kwargs:
kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
if 'Text Encoder' in shared.opts.bnb_quantization and 'text_encoder_3' not in kwargs:
kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
return kwargs
@@ -127,7 +125,7 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
kwargs = {}
kwargs = load_overrides(kwargs, cache_dir)
if fn is None or not os.path.exists(fn):
if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
kwargs = load_quants(kwargs, repo_id, cache_dir)
loader = diffusers.StableDiffusion3Pipeline.from_pretrained
modules/sd_models.py
+12 -6
View File
@@ -371,17 +371,19 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
class OffloadHook(accelerate.hooks.ModelHook):
def __init__(self):
def __init__(self, checkpoint_name):
if shared.opts.diffusers_offload_max_gpu_memory > 1:
shared.opts.diffusers_offload_max_gpu_memory = 0.75
if shared.opts.diffusers_offload_max_cpu_memory > 1:
shared.opts.diffusers_offload_max_cpu_memory = 0.75
self.checkpoint_name = checkpoint_name
self.min_watermark = shared.opts.diffusers_offload_min_gpu_memory
self.max_watermark = shared.opts.diffusers_offload_max_gpu_memory
self.cpu_watermark = shared.opts.diffusers_offload_max_cpu_memory
self.gpu = int(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024)
self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024)
self.offload_map = {}
self.param_map = {}
gpu = f'{shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory:.3f}-{shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory}:{shared.gpu_memory}'
shared.log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f}')
self.validate()
@@ -440,12 +442,12 @@ def apply_balanced_offload(sd_model, exclude=[]):
if sd_model.__class__.__name__ in excluded:
return sd_model
cached = True
if offload_hook_instance is None or offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory or offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory:
cached = False
offload_hook_instance = OffloadHook()
checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None
if checkpoint_name is None:
checkpoint_name = sd_model.__class__.__name__
if offload_hook_instance is None or offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory or offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory or checkpoint_name != offload_hook_instance.checkpoint_name:
cached = False
offload_hook_instance = OffloadHook(checkpoint_name)
def get_pipe_modules(pipe):
if hasattr(pipe, "_internal_dict"):
@@ -461,11 +463,13 @@ def apply_balanced_offload(sd_model, exclude=[]):
if not isinstance(module, torch.nn.Module):
continue
try:
module_size = sum(p.numel()*p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
except Exception as e:
shared.log.error(f'Offload: type=balanced op=calc module={module_name} {e}')
module_size = 0
offload_hook_instance.offload_map[module_name] = module_size
offload_hook_instance.param_map[module_name] = param_num
modules[module_name] = module_size
modules = sorted(modules.items(), key=lambda x: x[1], reverse=True)
return modules
@@ -494,7 +498,9 @@ def apply_balanced_offload(sd_model, exclude=[]):
if do_offload:
module = module.to(devices.cpu, non_blocking=True)
used_gpu -= module_size
debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} component={module.__class__.__name__} size={module_size:.3f}')
if not cached:
shared.log.debug(f'Offload: type=balanced module={module_name} cls={module.__class__.__name__} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}')
debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} module={module.__class__.__name__} size={module_size:.3f}')
except Exception as e:
if 'out of memory' in str(e):
devices.torch_gc(fast=True, force=True, reason='oom')
modules/textual_inversion/textual_inversion.py
+4 -1
View File
@@ -11,6 +11,7 @@ from modules.files_cache import directory_files, directory_mtime, extension_filt
debug = shared.log.trace if os.environ.get('SD_TI_DEBUG', None) is not None else lambda *args, **kwargs: None
debug('Trace: TEXTUAL INVERSION')
supported_models = ['ldm', 'sd', 'sdxl']
def list_embeddings(*dirs):
@@ -370,7 +371,7 @@ class EmbeddingDatabase:
self.skipped_embeddings[name] = embedding
def load_from_dir(self, embdir):
if sd_models.model_data.sd_model is None:
if not shared.sd_loaded:
shared.log.info('Skipping embeddings load: model not loaded')
return
if not os.path.isdir(embdir.path):
@@ -390,6 +391,8 @@ class EmbeddingDatabase:
def load_textual_inversion_embeddings(self, force_reload=False):
if not shared.sd_loaded:
return
if shared.sd_model_type not in supported_models:
return
t0 = time.time()
if not force_reload:
need_reload = False