pixart sigma improvements

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Pending
 
+- SC Lora
+- HunyuanDiT 1.1
 - Diffusers==0.30.0
 - https://github.com/huggingface/diffusers/issues/8546
 - https://github.com/huggingface/diffusers/pull/8566
@@ -12,9 +14,13 @@
 ### Highlights for 2024-06-17
 
 Following zero-day **SD3** release, a week later here's a refresh with more than a few improvements.  
-Other than SD3, highlight is (finally) new release of **Torch-DirectML** as well as ability to use explicit **T5** text-encoder in many supported models!  
+But there's more than SD3:
+- support for **PixArt-Sigma** in small/medium/large variants AND using 4/8/16bit quantized T5 text-encoder!  
+- (finally) new release of **Torch-DirectML**  
 
-### Improvements: SD3
+### Models
+
+#### Stable Diffusion 3
 
 - enable taesd preview and non-full quality mode  
 - enable base LoRA support  
@@ -28,10 +34,18 @@ Other than SD3, highlight is (finally) new release of **Torch-DirectML** as well
   higher shift means model will spend more time on structure and less on details  
 - add support for selecting text encoder in xyz grid
 
+#### Pixart-Σ
+
+- Add *small* (512px) and *large* (2k) variations, in addition to existing *medium* (1k)  
+- Add support for 4/8bit quantized t5 text encoder  
+  *note* by default pixart-Σ uses full fp16 t5 encoder with large memory footprint  
+  simply select in *settings -> model -> text encoder* before or after model load  
+
 ### Improvements: General
 
+- support FP4 quantized T5 text encoder, in addtion to existing FP8 and FP16
 - support for T5 text-encoder loader in **all** models that use T5  
-  *example*: load FP8 quantized T5 text-encoder into PixArt Sigma  
+  *example*: load FP4 or FP8 quantized T5 text-encoder into PixArt Sigma or Stable Cascade!
 - support for `torch-directml` **0.2.2**, thanks @lshqqytiger!  
   *note*: new directml is finally based on modern `torch` 2.3.1!  
 - extra networks: info display now contains link to source url if model if its known  

cli/api-json.py

@@ -45,10 +45,7 @@ if __name__ == "__main__":
     log.info(f'api-json: {args}')
     if os.path.isfile(args.json[0]):
         with open(args.json[0], 'r', encoding='ascii') as f:
-            txt = f.read()
-            txt = txt.encode('ascii')
-            print('HERE', txt)
-            dct = json.loads(txt)
+            dct = json.load(f) # TODO fails with b64 encoded images inside json due to string encoding
     else:
         dct = json.loads(args.json[0])
     res = post(endpoint=args.endpoint[0], payload=dct)

extensions-builtin/Lora/network_overrides.py

@@ -1,7 +1,7 @@
 from modules import shared
 
 
-maybe_diffusers = [
+maybe_diffusers = [ # forced if lora_maybe_diffusers is enabled
     'aaebf6360f7d', # sd15-lcm
     '3d18b05e4f56', # sdxl-lcm
     'b71dcb732467', # sdxl-tcd
@@ -19,14 +19,26 @@ maybe_diffusers = [
     '8cca3706050b', # hyper-sdxl-1step
 ]
 
-force_diffusers = [
+force_diffusers = [ # forced always
     '816d0eed49fd', # flash-sdxl
     'c2ec22757b46', # flash-sd15
 ]
 
+force_models = [ # forced always
+    'sd3',
+    'sc',
+    'hunyuandit',
+    'kandinsky',
+]
+
+force_classes = [ # forced always
+]
+
+
 def check_override(shorthash=''):
     force = False
-    force = force or (shared.sd_model_type == 'sd3') # TODO sd3 forced diffusers for lora load
+    force = force or (shared.sd_model_type in force_classes)
+    force = force or (shared.sd_model.__class__.__name__ in force_classes)
     if len(shorthash) < 4:
         return force
     force = force or (any(x.startswith(shorthash) for x in maybe_diffusers) if shared.opts.lora_maybe_diffusers else False)

html/reference.json

@@ -160,15 +160,30 @@
     "preview": "PixArt-alpha--PixArt-XL-2-1024-MS.jpg",
     "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
   },
-  "Pixart-Σ": {
-    "path": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
+  "Pixart-Σ Small": {
+    "path": "huggingface/PixArt-alpha/PixArt-Sigma-XL-2-512-MS",
     "desc": "PixArt-Σ, a Diffusion Transformer model (DiT) capable of directly generating images at 4K resolution. PixArt-Σ represents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts.",
     "preview": "PixArt-alpha--pixart_sigma_sdxlvae_T5_diffusers.jpg",
+    "skip": true,
+    "extras": "width: 512, height: 512, sampler: Default, cfg_scale: 2.0"
+  },
+  "Pixart-Σ Medium": {
+    "path": "huggingface/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
+    "desc": "PixArt-Σ, a Diffusion Transformer model (DiT) capable of directly generating images at 4K resolution. PixArt-Σ represents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts.",
+    "preview": "PixArt-alpha--pixart_sigma_sdxlvae_T5_diffusers.jpg",
+    "skip": true,
+    "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
+  },
+  "Pixart-Σ Large": {
+    "path": "huggingface/PixArt-alpha/PixArt-Sigma-XL-2-2K-MS",
+    "desc": "PixArt-Σ, a Diffusion Transformer model (DiT) capable of directly generating images at 4K resolution. PixArt-Σ represents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts.",
+    "preview": "PixArt-alpha--pixart_sigma_sdxlvae_T5_diffusers.jpg",
+    "skip": true,
     "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
   },
   
-  "Tencent HunyuanDiT": {
-    "path": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
+  "Tencent HunyuanDiT 1.1": {
+    "path": "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers",
     "desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
     "preview": "Tencent-Hunyuan-HunyuanDiT.jpg",
     "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"

javascript/ui.js

@@ -424,6 +424,7 @@ function selectVAE(name) {
 }
 
 function selectReference(name) {
+  log(`Select reference: ${name}`);
   desiredCheckpointName = name;
   gradioApp().getElementById('change_reference').click();
 }

modules/model_pixart.py

@@ -0,0 +1,29 @@
+import diffusers
+
+
+def load_pixart(checkpoint_info, diffusers_load_config={}):
+    from modules import shared, devices, modelloader, model_t5
+    modelloader.hf_login()
+    # shared.opts.data['cuda_dtype'] = 'FP32' # override
+    # shared.opts.data['diffusers_model_cpu_offload'] = True # override
+    # devices.set_cuda_params()
+    fn = checkpoint_info.path.replace('huggingface/', '')
+    t5 = model_t5.load_t5(shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+    transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
+        fn,
+        subfolder = 'transformer',
+        cache_dir = shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    transformer.to(devices.device)
+    kwargs = { 'transformer': transformer }
+    if t5 is not None:
+        kwargs['text_encoder'] = t5
+    pipe = diffusers.PixArtSigmaPipeline.from_pretrained(
+        'PixArt-alpha/PixArt-Sigma-XL-2-1024-MS',
+        cache_dir = shared.opts.diffusers_dir,
+        **kwargs,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe

modules/model_sd3.py

@@ -1,14 +1,7 @@
 import os
-import warnings
 import torch
 import diffusers
 import transformers
-import rich.traceback
-
-
-rich.traceback.install()
-warnings.filterwarnings(action="ignore", category=FutureWarning)
-loggedin = False
 
 
 def load_sd3(fn=None, cache_dir=None, config=None):
@@ -48,7 +41,7 @@ def load_sd3(fn=None, cache_dir=None, config=None):
                 ),
                 'text_encoder_3': None,
             }
-        elif fn_size < 1e10: # if model is below 10gb it does not have te4
+        elif fn_size < 1e10: # if model is below 10gb it does not have te3
             kwargs = {
                 'text_encoder_3': None,
             }
@@ -69,63 +62,3 @@ def load_sd3(fn=None, cache_dir=None, config=None):
     diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["stable-diffusion-3"] = diffusers.StableDiffusion3Img2ImgPipeline
     devices.torch_gc()
     return pipe
-
-
-def load_t5(pipe, module, te3=None, cache_dir=None):
-    from modules import devices, modelloader
-    repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
-    if pipe is None or not hasattr(pipe, module):
-        return pipe
-    if 'fp16' in te3.lower():
-        modelloader.hf_login()
-        t5 = transformers.T5EncoderModel.from_pretrained(
-            repo_id,
-            subfolder='text_encoder_3',
-            # torch_dtype=dtype,
-            cache_dir=cache_dir,
-            torch_dtype=pipe.text_encoder.dtype,
-        )
-        setattr(pipe, module, t5)
-    elif 'fp4' in te3.lower():
-        modelloader.hf_login()
-        from installer import install
-        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=pipe.text_encoder.dtype,
-        )
-        setattr(pipe, module, t5)
-    elif 'fp8' in te3.lower():
-        modelloader.hf_login()
-        from installer import install
-        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=pipe.text_encoder.dtype,
-        )
-        setattr(pipe, module, t5)
-        """
-        if hasattr(pipe, 'remove_all_hooks'):
-            pipe.remove_all_hooks()
-            nn = getattr(pipe, module)
-            import accelerate
-            accelerate.hooks.remove_hook_from_module(nn, recurse=True)
-            nn.to(device=devices.device)
-        """
-    else:
-        setattr(pipe, module, None)
-    if getattr(pipe, 'text_encoder_3', None) is not None and getattr(pipe, 'tokenizer_3', None) is None: # not needed anymore
-        pipe.tokenizer_3 = transformers.T5TokenizerFast.from_pretrained(
-            repo_id,
-            subfolder='tokenizer_3',
-            cache_dir=cache_dir,
-        )
-    devices.torch_gc()

modules/model_t5.py

@@ -0,0 +1,51 @@
+import transformers
+
+
+def load_t5(t5=None, cache_dir=None):
+    from modules import devices, modelloader
+    repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
+    if 'fp16' in t5.lower():
+        modelloader.hf_login()
+        t5 = transformers.T5EncoderModel.from_pretrained(
+            repo_id,
+            subfolder='text_encoder_3',
+            # torch_dtype=dtype,
+            cache_dir=cache_dir,
+            torch_dtype=devices.dtype,
+        )
+    elif 'fp4' in t5.lower():
+        modelloader.hf_login()
+        from installer import install
+        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():
+        modelloader.hf_login()
+        from installer import install
+        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,
+        )
+    else:
+        t5 = None
+    return t5
+
+
+def set_t5(pipe, module, t5=None, cache_dir=None):
+    from modules import devices
+    if pipe is None or not hasattr(pipe, module):
+        return pipe
+    t5 = load_t5(t5=t5, cache_dir=cache_dir)
+    setattr(pipe, module, t5)
+    devices.torch_gc()

modules/modelloader.py

@@ -204,7 +204,6 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
     shared.log.debug(f'Diffusers downloading: id="{hub_id}" args={download_config}')
     token = token or shared.opts.huggingface_token
     if token is not None and len(token) > 2:
-        shared.log.debug(f"Diffusers authentication: {token}")
         hf_login(token)
     pipeline_dir = None
 
@@ -318,6 +317,10 @@ def get_reference_opts(name: str, quiet=False):
         if k == name or model_name == name:
             model_opts = v
             break
+        model_name = model_name.replace('huggingface/', '')
+        if k == name or model_name == name:
+            model_opts = v
+            break
     if not model_opts:
         # shared.log.error(f'Reference: model="{name}" not found')
         return {}

modules/sd_models.py

@@ -202,11 +202,17 @@ def get_closet_checkpoint_match(search_string):
     if checkpoint_info is not None:
         return checkpoint_info
     found = sorted([info for info in checkpoints_list.values() if search_string in info.title], key=lambda x: len(x.title))
-    if found:
+    if found and len(found) > 0:
         return found[0]
     found = sorted([info for info in checkpoints_list.values() if search_string.split(' ')[0] in info.title], key=lambda x: len(x.title))
-    if found:
+    if found and len(found) > 0:
         return found[0]
+    for v in shared.reference_models.values():
+        if search_string in v['path'] or os.path.basename(search_string) in v['path']:
+            model_name = search_string.replace('huggingface/', '')
+            checkpoint_info = CheckpointInfo(v['path']) # create a virutal model info
+            checkpoint_info.type = 'huggingface'
+            return checkpoint_info
     return None
 
 
@@ -565,34 +571,20 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 # elif size < 0: # unknown
                 #    guess = 'Stable Diffusion 2B'
                 elif size >= 5791 and size <= 5799: # 5795
-                    if not shared.native:
-                        warn(f'Model detected as SD-XL refiner model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     if op == 'model':
                         warn(f'Model detected as SD-XL refiner model, but attempting to load a base model: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL Refiner'
                 elif (size >= 6611 and size <= 7220): # 6617, HassakuXL is 6776, monkrenRealisticINT_v10 is 7217
-                    if not shared.native:
-                        warn(f'Model detected as SD-XL base model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL'
                 elif size >= 3361 and size <= 3369: # 3368
-                    if not shared.native:
-                        warn(f'Model detected as SD upscale model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion Upscale'
                 elif size >= 4891 and size <= 4899: # 4897
-                    if not shared.native:
-                        warn(f'Model detected as SD XL inpaint model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL Inpaint'
                 elif size >= 9791 and size <= 9799: # 9794
-                    if not shared.native:
-                        warn(f'Model detected as SD XL instruct pix2pix model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL Instruct'
                 elif size > 3138 and size < 3142: #3140
-                    if not shared.native:
-                        warn(f'Model detected as Segmind Vega model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL'
                 elif size > 5692 and size < 5698 or size > 4134 and size < 4138:
-                    if not shared.native:
-                        warn(f'Model detected as Stable Diffusion 3 model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion 3'
             # guess by name
             """
@@ -602,34 +594,20 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 guess = 'Latent Consistency Model'
             """
             if 'instaflow' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as InstaFlow model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'InstaFlow'
             if 'segmoe' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as SegMoE model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'SegMoE'
             if 'hunyuandit' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as Tenecent HunyuanDiT model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'HunyuanDiT'
             if 'pixart-xl' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as PixArt Alpha model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'PixArt-Alpha'
             if 'stable-diffusion-3' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as Stable Diffusion 3 model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'Stable Diffusion 3'
             if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as Stable Cascade model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 if devices.dtype == torch.float16:
                     warn('Stable Cascade does not support Float16')
                 guess = 'Stable Cascade'
             if 'pixart-sigma' in f.lower():
-                if not shared.native:
-                    warn(f'Model detected as PixArt-Sigma model, but attempting to load using backend=original: {op}={f} size={size} MB')
                 guess = 'PixArt-Sigma'
             # switch for specific variant
             if guess == 'Stable Diffusion' and 'inpaint' in f.lower():
@@ -996,14 +974,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     return
             elif model_type in ['PixArt-Sigma']: # forced pipeline
                 try:
-                    # shared.opts.data['cuda_dtype'] = 'FP32' # override
-                    # shared.opts.data['diffusers_model_cpu_offload'] = True # override
-                    devices.set_cuda_params()
-                    sd_model = diffusers.PixArtSigmaPipeline.from_pretrained(
-                        checkpoint_info.path,
-                        use_safetensors=True,
-                        cache_dir=shared.opts.diffusers_dir,
-                        **diffusers_load_config)
+                    from modules.model_pixart import load_pixart
+                    sd_model = load_pixart(checkpoint_info, diffusers_load_config)
                 except Exception as e:
                     shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
                     if debug_load:
@@ -1533,13 +1505,13 @@ def reload_text_encoder(initial=False):
     signature = inspect.signature(shared.sd_model.__class__.__init__, follow_wrapped=True, eval_str=True).parameters
     t5 = [k for k, v in signature.items() if 'T5EncoderModel' in str(v)]
     if len(t5) > 0:
-        from modules.model_sd3 import load_t5
+        from modules.model_t5 import set_t5
         shared.log.debug(f'Load: t5={shared.opts.sd_text_encoder} module="{t5[0]}"')
-        load_t5(pipe=shared.sd_model, module=t5[0], te3=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+        set_t5(pipe=shared.sd_model, module=t5[0], t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
     elif hasattr(shared.sd_model, 'text_encoder_3'):
-        from modules.model_sd3 import load_t5
+        from modules.model_t5 import set_t5
         shared.log.debug(f'Load: t5={shared.opts.sd_text_encoder} module="text_encoder_3"')
-        load_t5(pipe=shared.sd_model, module='text_encoder_3', te3=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+        set_t5(pipe=shared.sd_model, module='text_encoder_3', t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
 
 
 def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', force=False):