From 2a27cd916d0524f900d9850be6be8987e56ac45d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 18 Jun 2024 13:21:13 -0400 Subject: [PATCH] pixart sigma improvements --- CHANGELOG.md | 20 +++++- cli/api-json.py | 5 +- extensions-builtin/Lora/network_overrides.py | 18 ++++- html/reference.json | 23 +++++-- javascript/ui.js | 1 + modules/model_pixart.py | 29 ++++++++ modules/model_sd3.py | 69 +------------------- modules/model_t5.py | 51 +++++++++++++++ modules/modelloader.py | 5 +- modules/sd_models.py | 56 ++++------------ 10 files changed, 152 insertions(+), 125 deletions(-) create mode 100644 modules/model_pixart.py create mode 100644 modules/model_t5.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 363020797..1167844ef 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/cli/api-json.py b/cli/api-json.py index 889b70d01..e8c5270fb 100755 --- a/cli/api-json.py +++ b/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) diff --git a/extensions-builtin/Lora/network_overrides.py b/extensions-builtin/Lora/network_overrides.py index 943575630..dacecd0bd 100644 --- a/extensions-builtin/Lora/network_overrides.py +++ b/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) diff --git a/html/reference.json b/html/reference.json index 2ffdba6a4..a3f62e7e1 100644 --- a/html/reference.json +++ b/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" diff --git a/javascript/ui.js b/javascript/ui.js index 8b4a5ce4a..c7ecd6680 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -424,6 +424,7 @@ function selectVAE(name) { } function selectReference(name) { + log(`Select reference: ${name}`); desiredCheckpointName = name; gradioApp().getElementById('change_reference').click(); } diff --git a/modules/model_pixart.py b/modules/model_pixart.py new file mode 100644 index 000000000..4ae7eb21c --- /dev/null +++ b/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 diff --git a/modules/model_sd3.py b/modules/model_sd3.py index 4d8aaf250..81470a97e 100644 --- a/modules/model_sd3.py +++ b/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() diff --git a/modules/model_t5.py b/modules/model_t5.py new file mode 100644 index 000000000..02764ea57 --- /dev/null +++ b/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() diff --git a/modules/modelloader.py b/modules/modelloader.py index 0045478a3..fd9e41edb 100644 --- a/modules/modelloader.py +++ b/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 {} diff --git a/modules/sd_models.py b/modules/sd_models.py index b06afefe8..ef18ebf5e 100644 --- a/modules/sd_models.py +++ b/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):