expermental t5 gguf support

modules/loader.py

@@ -38,6 +38,7 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
     torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
 timer.startup.record("torch")
 
+
 import transformers # pylint: disable=W0611,C0411
 timer.startup.record("transformers")
 

modules/model_flux.py

@@ -105,23 +105,22 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
         repo_path = checkpoint_info.path
     from installer import install
     install('bitsandbytes', quiet=True)
-    from diffusers import FluxTransformer2DModel
     quant = get_quant(repo_path)
     try:
         if quant == 'fp8':
             quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=devices.dtype)
             debug(f'Quantization: {quantization_config}')
-            transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+            transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
         elif quant == 'fp4':
             quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4')
             debug(f'Quantization: {quantization_config}')
-            transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+            transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
         elif quant == 'nf4':
             quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4')
             debug(f'Quantization: {quantization_config}')
-            transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+            transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
         else:
-            transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config)
+            transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config)
     except Exception as e:
         shared.log.error(f"Loading FLUX: Failed to load BnB transformer: {e}")
         transformer, text_encoder_2 = None, None
@@ -131,7 +130,24 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
     return transformer, text_encoder_2
 
 
+def load_flux_gguf(file_path): # TODO add support for GGUF flux models
+    shared.log.error(f"Loading FLUX: GGUF UNET is not supported: {file_path}")
+    """
+    with torch.device("meta"):
+        transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join("configs", "flux", "transformer", "config.json")).to(dtype=devices.dtype)
+    # from .modeling_gguf_pytorch_utils import load_gguf_checkpoint
+    from modules.model_te import install_gguf
+    install_gguf()
+    from transformers.modeling_gguf_pytorch_utils import load_gguf_checkpoint
+    state_dict = load_gguf_checkpoint(file_path, return_tensors=True)["tensors"]
+    return transformer, None
+    """
+    return None, None
+
+
 def load_transformer(file_path): # triggered by opts.sd_unet change
+    if file_path is None or not os.path.exists(file_path):
+        return
     transformer = None
     quant = get_quant(file_path)
     diffusers_load_config = {
@@ -140,7 +156,9 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         "cache_dir": shared.opts.hfcache_dir,
     }
     shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant={quant} dtype={devices.dtype}')
-    if quant == 'qint8' or quant == 'qint4':
+    if 'gguf' in file_path.lower():
+        _transformer, _text_encoder_2 = load_flux_gguf(file_path)
+    elif quant == 'qint8' or quant == 'qint4':
         _transformer, _text_encoder_2 = load_flux_quanto(file_path)
         if _transformer is not None:
             transformer = _transformer
@@ -154,8 +172,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         if _transformer is not None:
             transformer = _transformer
     else:
-        from diffusers import FluxTransformer2DModel
-        transformer = FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
+        transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
     if transformer is None:
         shared.log.error('Failed to load UNet model')
     return transformer
@@ -185,6 +202,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
                 sd_unet.failed_unet.append(shared.opts.sd_unet)
         except Exception as e:
             shared.log.error(f"Loading FLUX: Failed to load UNet: {e}")
+            shared.opts.sd_unet = 'None'
             if debug:
                 from modules import errors
                 errors.display(e, 'FLUX UNet:')
@@ -192,11 +210,12 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         try:
             debug(f'Loading FLUX: t5="{shared.opts.sd_text_encoder}"')
             from modules.model_te import load_t5
-            _text_encoder_2 = load_t5(t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+            _text_encoder_2 = load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
             if _text_encoder_2 is not None:
                 text_encoder_2 = _text_encoder_2
         except Exception as e:
             shared.log.error(f"Loading FLUX: Failed to load T5: {e}")
+            shared.opts.sd_text_encoder = 'None'
             if debug:
                 from modules import errors
                 errors.display(e, 'FLUX T5:')
@@ -211,6 +230,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
                 vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
         except Exception as e:
             shared.log.error(f"Loading FLUX: Failed to load VAE: {e}")
+            shared.opts.sd_vae = 'None'
             if debug:
                 from modules import errors
                 errors.display(e, 'FLUX VAE:')

modules/model_pixart.py

@@ -8,7 +8,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     # shared.opts.data['diffusers_offload_mode}'] = "model" # override
     # devices.set_cuda_params()
     fn = checkpoint_info.path.replace('huggingface/', '')
-    t5 = model_te.load_t5(shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+    t5 = model_te.load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
     transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
         fn,
         subfolder = 'transformer',

modules/model_te.py

@@ -12,12 +12,34 @@ debug = os.environ.get('SD_LOAD_DEBUG', None) is not None
 loaded_te = None
 
 
-def load_t5(t5=None, cache_dir=None):
+def install_gguf():
+    # pip install git+https://github.com/junejae/transformers@feature/t5-gguf
+    install('gguf', quiet=True)
+    # https://github.com/ggerganov/llama.cpp/issues/9566
+    import gguf
+    scripts_dir = os.path.join(os.path.dirname(gguf.__file__), '..', 'scripts')
+    if os.path.exists(scripts_dir):
+        os.rename(scripts_dir, scripts_dir + '_gguf')
+    # monkey patch transformers so they detect gguf pacakge correctly
+    import importlib
+    transformers.utils.import_utils._is_gguf_available = True # pylint: disable=protected-access
+    transformers.utils.import_utils._gguf_version = importlib.metadata.version('gguf') # pylint: disable=protected-access
+
+
+def load_t5(name=None, cache_dir=None):
+    global loaded_te # pylint: disable=global-statement
+    if name is None:
+        return
     from modules import modelloader
     modelloader.hf_login()
     repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
-    fn = te_dict.get(t5) if t5 in te_dict else None
-    if fn is not None and 'fp8' in t5.lower():
+    fn = te_dict.get(name) if name in te_dict else None
+    if fn is not None and 'gguf' in name.lower():
+        install_gguf()
+        with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f:
+            t5_config = transformers.T5Config(**json.load(f))
+        t5 = transformers.T5EncoderModel.from_pretrained(None, gguf_file=fn, config=t5_config, device_map="auto", cache_dir=cache_dir, torch_dtype=devices.dtype)
+    elif fn is not None and 'fp8' in name.lower():
         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))
@@ -42,22 +64,22 @@ def load_t5(t5=None, cache_dir=None):
             t5_config = transformers.T5Config(**json.load(f))
         state_dict = load_file(fn)
         t5 = transformers.T5EncoderModel.from_pretrained(None, state_dict=state_dict, config=t5_config)
-    elif 'fp16' in t5.lower():
+    elif 'fp16' in name.lower():
         t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
-    elif 'fp4' in t5.lower():
+    elif 'fp4' in name.lower():
         install('bitsandbytes', quiet=True)
         quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True)
         t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
-    elif 'fp8' in t5.lower():
+    elif 'fp8' in name.lower():
         install('bitsandbytes', quiet=True)
         quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True)
         t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
-    elif 'qint8' in t5.lower():
+    elif 'qint8' in name.lower():
         install('optimum-quanto', quiet=True)
         from modules.sd_models_compile import optimum_quanto_model
         t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
         t5 = optimum_quanto_model(t5, weights="qint8", activations="none")
-    elif 'int8' in t5.lower():
+    elif 'int8' in name.lower():
         install('nncf==2.7.0', quiet=True)
         from modules.sd_models_compile import nncf_compress_model
         from modules.sd_hijack import NNCF_T5DenseGatedActDense
@@ -70,6 +92,8 @@ def load_t5(t5=None, cache_dir=None):
         t5 = nncf_compress_model(t5)
     else:
         t5 = None
+    if t5 is not None:
+        loaded_te = name
     return t5
 
 
@@ -80,7 +104,7 @@ def set_t5(pipe, module, t5=None, cache_dir=None):
     if pipe is None or not hasattr(pipe, module):
         return pipe
     try:
-        t5 = load_t5(t5=t5, cache_dir=cache_dir)
+        t5 = load_t5(name=t5, cache_dir=cache_dir)
     except Exception as e:
         shared.log.error(f'Load module: type={module} class="T5" file="{shared.opts.sd_text_encoder}" {e}')
         if debug:
@@ -103,7 +127,7 @@ def set_t5(pipe, module, t5=None, cache_dir=None):
     return pipe
 
 
-def set_te(pipe):
+def set_clip(pipe):
     global loaded_te # pylint: disable=global-statement
     if loaded_te == shared.opts.sd_text_encoder:
         return
@@ -126,6 +150,7 @@ def set_te(pipe):
             import modules.prompt_parser_diffusers
             modules.prompt_parser_diffusers.cache.clear()
             move_model(pipe.text_encoder, devices.device)
+            devices.torch_gc()
     if 'vit-g' in shared.opts.sd_text_encoder.lower() and hasattr(shared.sd_model, 'text_encoder_2') and shared.sd_model.text_encoder_2.__class__.__name__ == 'CLIPTextModelWithProjection':
         try:
             config = transformers.PretrainedConfig.from_json_file('configs/sdxl/text_encoder_2/config.json')
@@ -144,11 +169,13 @@ def set_te(pipe):
             import modules.prompt_parser_diffusers
             modules.prompt_parser_diffusers.cache.clear()
             move_model(pipe.text_encoder_2, devices.device)
+            devices.torch_gc()
 
 
 def refresh_te_list():
     te_dict.clear()
-    for file in files_cache.list_files(shared.opts.te_dir, ext_filter=[".safetensors"]):
-        name = os.path.splitext(os.path.basename(file))[0]
+    for file in files_cache.list_files(shared.opts.te_dir, ext_filter=['.safetensors', '.gguf']):
+        basename = os.path.basename(file)
+        name = os.path.splitext(basename)[0] if '.safetensors' in basename else basename
         te_dict[name] = file
     shared.log.info(f'Available TEs: path="{shared.opts.te_dir}" items={len(te_dict)}')

modules/scripts.py

@@ -231,7 +231,7 @@ def list_scripts(scriptdirname, extension):
             else:
                 priority = priority + script.priority
             priority_list.append(ScriptFile(script.basedir, script.filename, script.path, priority))
-            debug(f'Adding script: {script.basedir} {script.filename} {script.path} {priority}')
+            debug(f'Adding script: folder="{script.basedir}" file="{script.filename}" full="{script.path}" priority={priority}')
     priority_sort = sorted(priority_list, key=lambda item: item.priority + item.path.lower(), reverse=False)
     return priority_sort
 
@@ -263,7 +263,7 @@ def load_scripts():
         for script_class in module.__dict__.values():
             if type(script_class) != type:
                 continue
-            debug(f'Registering script: {scriptfile.path}')
+            debug(f'Registering script: path="{scriptfile.path}"')
             if issubclass(script_class, Script):
                 scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))
             elif issubclass(script_class, scripts_postprocessing.ScriptPostprocessing):

modules/sd_models.py

@@ -1736,8 +1736,8 @@ def reload_text_encoder(initial=False):
         shared.log.debug(f'Load: t5={shared.opts.sd_text_encoder} module="text_encoder_3"')
         set_t5(pipe=shared.sd_model, module='text_encoder_3', t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
     elif hasattr(shared.sd_model, 'text_encoder') and 'vit' in shared.opts.sd_text_encoder.lower():
-        from modules.model_te import set_te
-        set_te(pipe=shared.sd_model)
+        from modules.model_te import set_clip
+        set_clip(pipe=shared.sd_model)
 
 
 def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', force=False):

modules/sd_unet.py

@@ -70,7 +70,8 @@ def load_unet(model):
 
 def refresh_unet_list():
     unet_dict.clear()
-    for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors"]):
-        name = os.path.splitext(os.path.basename(file))[0]
+    for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors", ".gguf"]):
+        basename = os.path.basename(file)
+        name = os.path.splitext(basename)[0] if ".safetensors" in basename else basename
         unet_dict[name] = file
     shared.log.info(f'Available UNets: path="{shared.opts.unet_dir}" items={len(unet_dict)}')