From b4df9a4de106f3ddd3cb098e00844e70d899af61 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 1 Sep 2024 22:56:15 -0400 Subject: [PATCH] jumbo update with flux.1 refactor, see changelog for details --- CHANGELOG.md | 46 ++++++-- modules/control/run.py | 28 +++-- modules/control/units/controlnet.py | 59 ++++++++--- modules/devices.py | 4 +- modules/model_flux.py | 159 ++++++++++++++++------------ modules/model_flux_nf4.py | 21 ++-- modules/model_t5.py | 23 +++- modules/sd_models.py | 18 +++- modules/sd_samplers.py | 1 + modules/sd_samplers_diffusers.py | 8 +- modules/sd_unet.py | 3 +- modules/shared.py | 1 + 12 files changed, 248 insertions(+), 123 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fd63d597..9d1e4507d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,22 +2,54 @@ ## Update for 2024-09-01 -- flux improve logging, warn when attempting to load unet as base model -- flux unet support fp8/fp4 quantization -- flux vae support fp16 -- flux lora support additional training tools (*1) -- flux model support loading all-in-one safetensors (*1) +Major refactor of FLUX.1 support: +- allow configuration of individual FLUX.1 model components: *transformer, text-encoder, vae* + model load will load selected components first and then initialize model using pre-loaded components + components that were not pre-loaded will be downloaded and initialized as needed + as usual, components can also be loaded after initial model load + *note*: use of transformer/unet is recommended as those are flux.1 finetunes + *note*: manually selecting vae and text-encoder is not recommended + *note*: mix-and-match of different quantizations for different components can lead to unexpected errors + - transformer/unet is list of manually downloaded safetensors + - vae is list of manually downloaded safetensors + - text-encoder is list of predefined and manually downloaded text-encoders +- model support loading all-in-one safetensors (*1) not recommended due to massive duplication of components, but added due to popular demand + each such model is 20-32GB in size vs ~11GB for typical unet fine-tune +- improve logging, warn when attempting to load unet as base model +- transformer/unet support fp8/fp4 quantization +- vae support fp16 (*1) +- lora support additional training tools (*1) +- support fuse-qkv projections (*1) + can speed up generate + enable via *settings -> compute -> fused projections* + +Other improvements: - taesd configurable number of layers can be used to speed-up taesd decoding by reducing number of ops e.g. if generating 1024px image, reducing layers by 1 will result in preview being 512px set via *settings -> live preview -> taesd decode layers* - xhinker prompt parser handle offloaded models -- t5 enum manually downloaded models (*2) +- control better handle offloading +- speed up some garbage collection ops +- sampler settings add dynamic shift + used by flow-matching samplers to adjust between structure and details +- sampler settings force base shift + improves quality of the flow-matching samplers +- t5 support manually downloaded models + applies to all models that use t5 transformer + +Work-in-progress: +- flux controlnet support: (*1) + - instanx models + - shakker-labs models + - TBD: add controlnet_mode for union models + - TBD: validate control_image vs input_type + - TBD: not enough values to unpack + - TBD: xlabs models *notes*: - (*1) requires `diffusers==0.31.0.dev0` -- (*2) work-in-progress ## Update for 2024-08-31 diff --git a/modules/control/run.py b/modules/control/run.py index 41204fa49..f978dbf36 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -544,17 +544,23 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini if is_generator: yield terminate('Control: attempting reference mode but image is none') return [], '', '', 'Reference mode without image' - elif unit_type == 'controlnet' and input_type == 1 and has_models: # Init image same as control - p.task_args['control_image'] = p.init_images # switch image and control_image - p.task_args['strength'] = p.denoising_strength - p.init_images = [p.override or input_image] * len(active_model) - elif unit_type == 'controlnet' and input_type == 2 and has_models: # Separate init image - if init_image is None: - shared.log.warning('Control: separate init image not provided') - init_image = input_image - p.task_args['control_image'] = p.init_images # switch image and control_image - p.task_args['strength'] = p.denoising_strength - p.init_images = [init_image] * len(active_model) + elif unit_type == 'controlnet' and has_models: + if input_type == 0: # Control only + if shared.sd_model_type == 'f1': + if is_generator: + yield terminate('Control: Flux control invalid input type') + return [], '', '', 'Flux control invalid input type' + elif input_type == 1: # Init image same as control + p.task_args['control_image'] = p.init_images # switch image and control_image + p.task_args['strength'] = p.denoising_strength + p.init_images = [p.override or input_image] * len(active_model) + elif input_type == 2: # Separate init image + if init_image is None: + shared.log.warning('Control: separate init image not provided') + init_image = input_image + p.task_args['control_image'] = p.init_images # switch image and control_image + p.task_args['strength'] = p.denoising_strength + p.init_images = [init_image] * len(active_model) if is_generator: image_txt = f'{blended_image.width}x{blended_image.height}' if blended_image is not None else 'None' diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index 419e119fc..a200f59a2 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -1,10 +1,10 @@ import os import time from typing import Union -from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline +from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline, ControlNetModel from modules.control.units import detect from modules.shared import log, opts, listdir -from modules import errors, sd_models +from modules import errors, sd_models, devices what = 'ControlNet' @@ -80,6 +80,7 @@ models = {} all_models = {} all_models.update(predefined_sd15) all_models.update(predefined_sdxl) +all_models.update(predefined_f1) cache_dir = 'models/control/controlnet' @@ -139,6 +140,19 @@ class ControlNet(): self.model = None self.model_id = None + def get_class(self): + import modules.shared + if modules.shared.sd_model_type == 'sd': + from diffusers import ControlNetModel as model_class # pylint: disable=reimported + elif modules.shared.sd_model_type == 'sdxl': + from diffusers import ControlNetModel as model_class # pylint: disable=reimported # sdxl shares same model class + elif modules.shared.sd_model_type == 'f1': + from diffusers import FluxControlNetModel as model_class + else: + log.error(f'Control {what}: type={modules.shared.sd_model_type} unsupported model') + return None + return model_class + def load_safetensors(self, model_path): name = os.path.splitext(model_path)[0] config_path = None @@ -164,7 +178,8 @@ class ControlNet(): config_path = f'{name}.json' if config_path is not None: self.load_config['original_config_file '] = config_path - self.model = ControlNetModel.from_single_file(model_path, **self.load_config) + cls = self.get_class() + self.model = cls.from_single_file(model_path, **self.load_config) def load(self, model_id: str = None) -> str: try: @@ -189,7 +204,8 @@ class ControlNet(): if '/bin' in model_path: model_path = model_path.replace('/bin', '') self.load_config['use_safetensors'] = False - self.model = ControlNetModel.from_pretrained(model_path, **self.load_config) + cls = self.get_class() + self.model = cls.from_pretrained(model_path, **self.load_config) if self.dtype is not None: self.model.to(self.dtype) if "ControlNet" in opts.nncf_compress_weights: @@ -223,7 +239,7 @@ class ControlNet(): class ControlNetPipeline(): - def __init__(self, controlnet: Union[ControlNetModel, list[ControlNetModel]], pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline], dtype = None): + def __init__(self, controlnet: Union[ControlNetModel, list[ControlNetModel]], pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline, FluxPipeline], dtype = None): t0 = time.time() self.orig_pipeline = pipeline self.pipeline = None @@ -231,6 +247,7 @@ class ControlNetPipeline(): log.error('Control model pipeline: model not loaded') return elif detect.is_sdxl(pipeline): + from diffusers import StableDiffusionXLControlNetPipeline self.pipeline = StableDiffusionXLControlNetPipeline( vae=pipeline.vae, text_encoder=pipeline.text_encoder, @@ -242,8 +259,8 @@ class ControlNetPipeline(): feature_extractor=getattr(pipeline, 'feature_extractor', None), controlnet=controlnet, # can be a list ) - sd_models.move_model(self.pipeline, pipeline.device) elif detect.is_sd15(pipeline): + from diffusers import StableDiffusionControlNetPipeline self.pipeline = StableDiffusionControlNetPipeline( vae=pipeline.vae, text_encoder=pipeline.text_encoder, @@ -257,17 +274,33 @@ class ControlNetPipeline(): ) sd_models.move_model(self.pipeline, pipeline.device) elif detect.is_f1(pipeline): - log.warning('Control model pipeline: class=FluxPipeline unsupported model type') + from diffusers import FluxControlNetPipeline + self.pipeline = FluxControlNetPipeline( + vae=pipeline.vae, + text_encoder=pipeline.text_encoder, + text_encoder_2=pipeline.text_encoder_2, + tokenizer=pipeline.tokenizer, + tokenizer_2=pipeline.tokenizer_2, + transformer=pipeline.transformer, + scheduler=pipeline.scheduler, + controlnet=controlnet, # can be a list + ) else: log.error(f'Control {what} pipeline: class={pipeline.__class__.__name__} unsupported model type') return - if dtype is not None and self.pipeline is not None: - self.pipeline = self.pipeline.to(dtype) - t1 = time.time() - if self.pipeline is not None: - log.debug(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}') - else: + + if self.pipeline is None: log.error(f'Control {what} pipeline: not initialized') + return + if dtype is not None: + self.pipeline = self.pipeline.to(dtype) + if opts.diffusers_offload_mode == 'none': + sd_models.move_model(self.pipeline, devices.device) + from modules.sd_models import set_diffuser_offload + set_diffuser_offload(self.pipeline, 'model') + + t1 = time.time() + log.debug(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}') def restore(self): self.pipeline = None diff --git a/modules/devices.py b/modules/devices.py index 1d6ae3703..37caa7f82 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -128,7 +128,7 @@ def get_device_for(task): return get_optimal_device() -def torch_gc(force=False): +def torch_gc(force=False, fast=False): t0 = time.time() mem = memstats.memory_stats() gpu = mem.get('gpu', {}) @@ -151,7 +151,7 @@ def torch_gc(force=False): return # actual gc - collected = gc.collect() # python gc + collected = gc.collect() if not fast else 0 # python gc if cuda_ok: try: with torch.cuda.device(get_cuda_device_string()): diff --git a/modules/model_flux.py b/modules/model_flux.py index 93cf1a641..29f8afaaa 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -5,11 +5,10 @@ import diffusers import transformers from safetensors.torch import load_file from huggingface_hub import hf_hub_download -from accelerate.utils import compute_module_sizes from modules import shared, devices -debug = os.environ.get('SD_LOAD_DEBUG', None) is not None +debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None def get_quant(file_path): @@ -26,7 +25,7 @@ def get_quant(file_path): return 'none' -def load_flux_quanto(checkpoint_info, diffusers_load_config, transformer_only=False): +def load_flux_quanto(checkpoint_info, diffusers_load_config, transformer, text_encoder_2): from installer import install install('optimum-quanto', quiet=True) try: @@ -41,51 +40,51 @@ def load_flux_quanto(checkpoint_info, diffusers_load_config, transformer_only=Fa repo_path = checkpoint_info else: repo_path = checkpoint_info.path - quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") - if not os.path.exists(quantization_map): - repo_id = checkpoint_info.name.replace('Diffusers/', '') - quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', **diffusers_load_config) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) - dtype = state_dict['context_embedder.bias'].dtype - with torch.device("meta"): - transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) - requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) - transformer.eval() - if transformer.dtype != devices.dtype: - try: - transformer = transformer.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"FLUX: Failed to cast transformer to {devices.dtype}, set dtype to {transformer.dtype}") - raise - if transformer_only: - return transformer, None - quantization_map = os.path.join(repo_path, "text_encoder_2", "quantization_map.json") - if not os.path.exists(quantization_map): - repo_id = checkpoint_info.name.replace('Diffusers/', '') - quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', **diffusers_load_config) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: - t5_config = transformers.T5Config(**json.load(f)) - state_dict = load_file(os.path.join(repo_path, "text_encoder_2", "model.safetensors")) - dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype - with torch.device("meta"): - text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) - requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) - text_encoder_2.eval() - if text_encoder_2.dtype != devices.dtype: - try: - text_encoder_2 = text_encoder_2.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"FLUX: Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2.dtype}") - raise - return transformer, text_encoder_2 + if transformer is None: + quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") + if not os.path.exists(quantization_map): + repo_id = checkpoint_info.name.replace('Diffusers/', '') + quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', **diffusers_load_config) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) + dtype = state_dict['context_embedder.bias'].dtype + with torch.device("meta"): + transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) + requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) + transformer.eval() + if transformer.dtype != devices.dtype: + try: + transformer = transformer.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"FLUX: Failed to cast transformer to {devices.dtype}, set dtype to {transformer.dtype}") + raise + + if text_encoder_2 is None: + quantization_map = os.path.join(repo_path, "text_encoder_2", "quantization_map.json") + if not os.path.exists(quantization_map): + repo_id = checkpoint_info.name.replace('Diffusers/', '') + quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', **diffusers_load_config) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: + t5_config = transformers.T5Config(**json.load(f)) + state_dict = load_file(os.path.join(repo_path, "text_encoder_2", "model.safetensors")) + dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype + with torch.device("meta"): + text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) + requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) + text_encoder_2.eval() + if text_encoder_2.dtype != devices.dtype: + try: + text_encoder_2 = text_encoder_2.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"FLUX: Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2.dtype}") + raise -def load_flux_bnb(checkpoint_info, diffusers_load_config, transformer_only=False): +def load_flux_bnb(checkpoint_info, diffusers_load_config, transformer, text_encoder_2): # pylint: disable=unused-argument if isinstance(checkpoint_info, str): repo_path = checkpoint_info else: @@ -96,18 +95,18 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config, transformer_only=False quant = get_quant(repo_path) if quant == 'fp8': quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True) - transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + if transformer is None: + transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) elif quant == 'fp4': quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True) - transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + if transformer is None: + transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) else: - transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) - if transformer_only: - return transformer, None - # TODO load text_encoder_2 + if transformer is None: + transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) -def load_transformer(file_path): # triggered by opts.sd_unet change +def load_transformer(file_path, transformer): # triggered by opts.sd_unet change quant = get_quant(file_path) diffusers_load_config = { "low_cpu_mem_usage": True, @@ -117,32 +116,60 @@ def load_transformer(file_path): # triggered by opts.sd_unet change shared.log.info(f'Loading UNet: type=FLUX file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant={quant} dtype={devices.dtype}') if 'nf4' in quant: from modules.model_flux_nf4 import load_flux_nf4 - transformer = load_flux_nf4(file_path, diffusers_load_config, transformer_only=True) + load_flux_nf4(file_path, diffusers_load_config, transformer, text_encoder_2='skip') elif quant == 'qint8' or quant == 'qint4': - transformer, _ = load_flux_quanto(file_path, diffusers_load_config, transformer_only=True) + load_flux_quanto(file_path, diffusers_load_config, transformer, text_encoder_2='skip') elif quant == 'fp8' or quant == 'fp4': - transformer, _ = load_flux_bnb(file_path, diffusers_load_config, transformer_only=True) + load_flux_bnb(file_path, diffusers_load_config, transformer, text_encoder_2='skip') else: from diffusers import FluxTransformer2DModel transformer = FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config) if transformer is None: shared.log.error('Failed to load UNet model') - if debug: - shared.log.debug(f'FLUX transformer: size={round(compute_module_sizes(transformer)[""] / 1024 / 1204)}') return transformer def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change quant = get_quant(checkpoint_info.path) - shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant={quant} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}') + shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" unet="{shared.opts.sd_unet}" t5="{shared.opts.sd_text_encoder}" vae="{shared.opts.sd_vae}" quant={quant} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}') + + transformer = None + text_encoder_2 = None + vae = None + + # load overrides if any + if shared.opts.sd_unet != 'None': + debug(f'Loading FLUX: unet="{shared.opts.sd_unet}"') + from modules import sd_unet + load_transformer(sd_unet.unet_dict[shared.opts.sd_unet], transformer) + if shared.opts.sd_text_encoder != 'None': + debug(f'Loading FLUX: t5="{shared.opts.sd_text_encoder}"') + from modules.model_t5 import load_t5 + text_encoder_2 = load_t5(t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir) + if shared.opts.sd_vae != 'None': + debug(f'Loading FLUX: vae="{shared.opts.sd_vae}"') + from modules import sd_vae + # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') + vae_file = sd_vae.vae_dict[shared.opts.sd_vae] + if os.path.exists(vae_file): + vae_config = os.path.join('configs', 'flux', 'vae', 'config.json') + vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config) + + # load quantized components if any if quant == 'nf4': from modules.model_flux_nf4 import load_flux_nf4 - pipe = load_flux_nf4(checkpoint_info, diffusers_load_config) - elif quant == 'qint8' or quant == 'qint4': - pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, transformer=None, text_encoder_2=None, **diffusers_load_config) - pipe.transformer, pipe.text_encoder_2 = load_flux_quanto(checkpoint_info, diffusers_load_config) - else: - pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) - if debug: - shared.log.debug(f'FLUX transformer: size={round(compute_module_sizes(pipe.transformer)[""] / 1024 / 1204)}') + load_flux_nf4(checkpoint_info, diffusers_load_config, transformer, text_encoder_2) + if quant == 'qint8' or quant == 'qint4': + load_flux_quanto(checkpoint_info, diffusers_load_config, transformer, text_encoder_2) + + # initialize pipeline with pre-loaded components + components = {} + if transformer is not None: + components['transformer'] = transformer + if text_encoder_2 is not None: + components['text_encoder_2'] = text_encoder_2 + if vae is not None: + components['vae'] = vae + debug(f'Loading FLUX: preloaded={list(components)}') + pipe = diffusers.FluxPipeline.from_pretrained('black-forest-labs/flux.1-dev', cache_dir=shared.opts.diffusers_dir, **components, **diffusers_load_config) return pipe diff --git a/modules/model_flux_nf4.py b/modules/model_flux_nf4.py index 61d852862..9e995b7f8 100644 --- a/modules/model_flux_nf4.py +++ b/modules/model_flux_nf4.py @@ -9,7 +9,6 @@ from transformers.quantizers.quantizers_utils import get_module_from_name from huggingface_hub import hf_hub_download from accelerate import init_empty_weights from accelerate.utils import set_module_tensor_to_device -from diffusers import FluxTransformer2DModel, FluxPipeline from diffusers.loaders.single_file_utils import convert_flux_transformer_checkpoint_to_diffusers import safetensors.torch from modules import shared, devices @@ -163,7 +162,7 @@ def create_quantized_param( module._parameters[tensor_name] = new_value # pylint: disable=protected-access -def load_flux_nf4(checkpoint_info, diffusers_load_config, transformer_only=False): +def load_flux_nf4(checkpoint_info, diffusers_load_config, transformer, text_encoder_2): load_bnb() if isinstance(checkpoint_info, str): repo_path = checkpoint_info @@ -190,11 +189,12 @@ def load_flux_nf4(checkpoint_info, diffusers_load_config, transformer_only=False converted_state_dict = original_state_dict with init_empty_weights(): + from diffusers import FluxTransformer2DModel config = FluxTransformer2DModel.load_config("black-forest-labs/flux.1-dev", subfolder="transformer") - model = FluxTransformer2DModel.from_config(config).to(devices.dtype) - expected_state_dict_keys = list(model.state_dict().keys()) + transformer = FluxTransformer2DModel.from_config(config).to(devices.dtype) + expected_state_dict_keys = list(transformer.state_dict().keys()) - _replace_with_bnb_linear(model, "nf4") + _replace_with_bnb_linear(transformer, "nf4") for param_name, param in converted_state_dict.items(): if param_name not in expected_state_dict_keys: @@ -202,15 +202,10 @@ def load_flux_nf4(checkpoint_info, diffusers_load_config, transformer_only=False is_param_float8_e4m3fn = hasattr(torch, "float8_e4m3fn") and param.dtype == torch.float8_e4m3fn if torch.is_floating_point(param) and not is_param_float8_e4m3fn: param = param.to(devices.dtype) - if not check_quantized_param(model, param_name): - set_module_tensor_to_device(model, param_name, device=0, value=param) + if not check_quantized_param(transformer, param_name): + set_module_tensor_to_device(transformer, param_name, device=0, value=param) else: - create_quantized_param(model, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True) + create_quantized_param(transformer, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True) del original_state_dict devices.torch_gc(force=True) - if transformer_only: - return model - else: - pipe = FluxPipeline.from_pretrained("black-forest-labs/flux.1-dev", transformer=model, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) - return pipe diff --git a/modules/model_t5.py b/modules/model_t5.py index fda3bc6fe..6a1f15dd5 100644 --- a/modules/model_t5.py +++ b/modules/model_t5.py @@ -1,6 +1,8 @@ import os +import json import torch import transformers +from safetensors.torch import load_file from modules import shared, devices, files_cache @@ -12,8 +14,25 @@ def load_t5(t5=None, cache_dir=None): repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers' fn = t5_dict.get(t5) if t5 in t5_dict else None if fn is not None: - shared.log.error(f'Loading T5: file="{fn}" unsupported') - t5 = None + from accelerate.utils import set_module_tensor_to_device + with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f: + t5_config = transformers.T5Config(**json.load(f)) + state_dict = load_file(fn) + dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype + with torch.device("meta"): + t5 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) + for param_name, param in state_dict.items(): + is_param_float8_e4m3fn = hasattr(torch, "float8_e4m3fn") and param.dtype == torch.float8_e4m3fn + if torch.is_floating_point(param) and not is_param_float8_e4m3fn: + param = param.to(devices.dtype) + set_module_tensor_to_device(t5, param_name, device=0, value=param) + t5.eval() + if t5.dtype != devices.dtype: + try: + t5 = t5.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"FLUX: Failed to cast text encoder to {devices.dtype}, set dtype to {t5.dtype}") + raise elif 'fp16' in t5.lower(): modelloader.hf_login() t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) diff --git a/modules/sd_models.py b/modules/sd_models.py index 5b1b9f021..c5e03010e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -713,6 +713,12 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): shared.log.debug(f'Setting {op}: enable fused projections') except Exception as e: shared.log.error(f'Error enabling fused projections: {e}') + if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'fuse_qkv_projections'): + try: + sd_model.transformer.fuse_qkv_projections() + shared.log.debug(f'Setting {op}: enable fused projections') + except Exception as e: + shared.log.error(f'Error enabling fused projections: {e}') if shared.opts.diffusers_eval: def eval_model(model, op=None, sd_model=None): # pylint: disable=unused-argument if hasattr(model, "requires_grad_"): @@ -734,6 +740,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'): if sd_model is None: shared.log.warning(f'{op} is not loaded') return + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if hasattr(sd_model, "enable_model_cpu_offload"): @@ -828,7 +835,7 @@ def apply_balanced_offload(sd_model): module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access except Exception as e: shared.log.error(f'Balanced offload: module={module_name} {e}') - devices.torch_gc() + devices.torch_gc(fast=True) apply_balanced_offload_to_module(sd_model) if hasattr(sd_model, "prior_pipe"): @@ -1046,16 +1053,17 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No unload_model_weights(op=op) return + shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"') + pipeline, model_type = detect_pipeline(checkpoint_info.path, op) + vae = None sd_vae.loaded_vae_file = None - if op == 'model' or op == 'refiner': + if model_type.startswith('Stable Diffusion') and (op == 'model' or op == 'refiner'): # preload vae for sd models vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename) vae = sd_vae.load_vae_diffusers(checkpoint_info.path, vae_file, vae_source) if vae is not None: diffusers_load_config["vae"] = vae - shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"') - pipeline, model_type = detect_pipeline(checkpoint_info.path, op) if os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface' or checkpoint_info.type == 'transformer': files = shared.walk_files(checkpoint_info.path, ['.safetensors', '.bin', '.ckpt']) if 'variant' not in diffusers_load_config and any('diffusion_pytorch_model.fp16' in f for f in files): # deal with diffusers lack of variant fallback when loading @@ -1462,7 +1470,7 @@ def set_diffuser_pipe(pipe, new_pipe_type): return pipe # skip specific pipelines - if n in ['StableDiffusionReferencePipeline', 'StableDiffusionAdapterPipeline', 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline']: + if n in ['StableDiffusionReferencePipeline', 'StableDiffusionAdapterPipeline', 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline', 'FluxPipeline', 'FluxControlNetPipeline']: # TODO flux does not have inpaint/img2img yet return pipe if 'Onnx' in pipe.__class__.__name__: return pipe diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 04d3e9d37..6d86b1bf1 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -72,6 +72,7 @@ def create_sampler(name, model): if shared.sd_model_type == 'f1': if 'base_image_seq_len' not in sampler.sampler.config or 'max_image_seq_len' not in sampler.sampler.config or 'base_shift' not in sampler.sampler.config or 'max_shift' not in sampler.sampler.config: shared.log.warning(f'FLUX: sampler="{name}" non compatible') + # sampler.sampler.register_to_config(base_image_seq_len=256, max_image_seq_len=4096, base_shift=0.5, max_shift=1.15) return None if not hasattr(model, 'scheduler_config'): model.scheduler_config = sampler.sampler.config.copy() diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 24f3d033e..a2e8adc43 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -68,8 +68,8 @@ config = { 'Euler EDM': { }, 'DPM++ 2M EDM': { 'solver_order': 2, 'solver_type': 'midpoint', 'final_sigmas_type': 'zero', 'algorithm_type': 'dpmsolver++' }, 'CMSI': { }, #{ 'sigma_min': 0.002, 'sigma_max': 80.0, 'sigma_data': 0.5, 's_noise': 1.0, 'rho': 7.0, 'clip_denoised': True }, - 'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, }, - 'Heun FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, }, + 'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, 'use_dynamic_shifting': False }, + 'Heun FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1 }, 'IPNDM': { }, } @@ -157,8 +157,10 @@ class DiffusionSampler: self.config['beta_start'] = shared.opts.schedulers_beta_start if 'beta_end' in self.config and shared.opts.schedulers_beta_end > 0: self.config['beta_end'] = shared.opts.schedulers_beta_end - if 'shift' in self.config and shared.opts.schedulers_shift != 1: + if 'shift' in self.config: self.config['shift'] = shared.opts.schedulers_shift + if 'use_dynamic_shifting' in self.config: + self.config['use_dynamic_shifting'] = shared.opts.schedulers_dynamic_shift if 'rescale_betas_zero_snr' in self.config: self.config['rescale_betas_zero_snr'] = shared.opts.schedulers_rescale_betas if 'timestep_spacing' in self.config and shared.opts.schedulers_timestep_spacing != 'default' and shared.opts.schedulers_timestep_spacing is not None: diff --git a/modules/sd_unet.py b/modules/sd_unet.py index 16d942a22..7271ce50a 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -30,7 +30,8 @@ def load_unet(model): model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype) if "Flux" in model.__class__.__name__: from modules.model_flux import load_transformer - transformer = load_transformer(unet_dict[shared.opts.sd_unet]) + transformer = None + load_transformer(unet_dict[shared.opts.sd_unet], transformer) if transformer is not None: model.transformer = None if shared.opts.diffusers_offload_mode == 'none': diff --git a/modules/shared.py b/modules/shared.py index 0fceaa7da..7be446482 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -742,6 +742,7 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"), 'schedulers_timesteps': OptionInfo('', "Timesteps"), "schedulers_rescale_betas": OptionInfo(False, "Rescale betas with zero terminal SNR", gr.Checkbox), 'schedulers_shift': OptionInfo(1, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1}), + 'schedulers_dynamic_shift': OptionInfo(True, "Sampler dynamic shift"), # managed from ui.py for backend original k-diffusion "schedulers_sep_kdiffusers": OptionInfo("

K-Diffusion specific config

", "", gr.HTML),