jumbo update with flux.1 refactor, see changelog for details

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
 

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'

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

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()):

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

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

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)

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

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()

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:

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':

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("<h2>K-Diffusion specific config</h2>", "", gr.HTML),