diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index ef53fb763..419e119fc 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -67,7 +67,14 @@ predefined_sdxl = { # 'StabilityAI Sketch R256': 'stabilityai/control-lora/control-LoRAs-rank256/control-lora-sketch-rank256.safetensors', } predefined_f1 = { - 'Shakker-Labs ControlNet Union': 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro', + "InstantX Union": 'InstantX/FLUX.1-dev-Controlnet-Union', + "InstantX Canny": 'InstantX/FLUX.1-dev-Controlnet-Canny', + "Shakker-Labs Union": 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro', + "Shakker-Labs Pose": 'Shakker-Labs/FLUX.1-dev-ControlNet-Pose', + "Shakker-Labs Depth": 'Shakker-Labs/FLUX.1-dev-ControlNet-Depth', + "XLabs-AI Canny": 'XLabs-AI/flux-controlnet-canny-v3', + "XLabs-AI Depth": 'XLabs-AI/flux-controlnet-depth-v3', + "XLabs-AI HED": 'XLabs-AI/flux-controlnet-hed-v3', } models = {} all_models = {} diff --git a/modules/model_flux.py b/modules/model_flux.py index 82a7a6456..4940a659e 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -1,32 +1,43 @@ +import os import json import torch import diffusers import transformers from safetensors.torch import load_file -from accelerate.utils import compute_module_sizes +from huggingface_hub import hf_hub_download from modules import shared, devices -def load_quanto_transformer(repo_path): +def load_quanto_transformer(checkpoint_info): from optimum.quanto import requantize # pylint: disable=no-name-in-module - with open(repo_path + "/" + "transformer/quantization_map.json", "r", encoding='utf8') as f: + 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', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: quantization_map = json.load(f) - state_dict = load_file(repo_path + "/" + "transformer/diffusion_pytorch_model.safetensors") + 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(repo_path + "/" + "transformer/config.json").to(dtype=dtype) + 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() return transformer -def load_quanto_text_encoder_2(repo_path): +def load_quanto_text_encoder_2(checkpoint_info): from optimum.quanto import requantize # pylint: disable=no-name-in-module - with open(repo_path + "/" + "text_encoder_2/quantization_map.json", "r", encoding='utf8') as f: + repo_path = checkpoint_info.path + 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', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: quantization_map = json.load(f) - with open(repo_path + "/" + "text_encoder_2/config.json", encoding='utf8') as 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(repo_path + "/" + "text_encoder_2/model.safetensors") + 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) @@ -78,8 +89,8 @@ def load_flux(checkpoint_info, diffusers_load_config): raise quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs) 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 = load_quanto_transformer(checkpoint_info.path) - pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info.path) + pipe.transformer = load_quanto_transformer(checkpoint_info) + pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info) if pipe.transformer.dtype != devices.dtype: try: pipe.transformer = pipe.transformer.to(dtype=devices.dtype) @@ -97,5 +108,6 @@ def load_flux(checkpoint_info, diffusers_load_config): if devices.dtype == torch.float16 and not shared.opts.no_half_vae: shared.log.warning("FLUX: does not support FP16 VAE, enabling no-half-vae") shared.opts.no_half_vae = True - shared.log.debug(f'FLUX computed size: {round(compute_module_sizes(pipe.transformer)[""] / 1024 / 1204)}') + # from accelerate.utils import compute_module_sizes + # shared.log.debug(f'FLUX computed size: {round(compute_module_sizes(pipe.transformer)[""] / 1024 / 1204)}') return pipe diff --git a/modules/model_stablecascade.py b/modules/model_stablecascade.py index 831c313fc..57c0d5678 100644 --- a/modules/model_stablecascade.py +++ b/modules/model_stablecascade.py @@ -4,6 +4,7 @@ import torch import diffusers from modules import shared, devices, sd_models + def get_timestep_ratio_conditioning(t, alphas_cumprod): s = torch.tensor([0.008]) # diffusers uses 0.003 while the original is 0.008 clamp_range = [0, 1] @@ -14,6 +15,7 @@ def get_timestep_ratio_conditioning(t, alphas_cumprod): ratio = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s return ratio + def load_text_encoder(path): from transformers import CLIPTextConfig, CLIPTextModelWithProjection from accelerate.utils.modeling import set_module_tensor_to_device @@ -131,9 +133,9 @@ def load_cascade_combined(checkpoint_info, diffusers_load_config): sd_model = StableCascadeCombinedPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) shared.log.debug(f'StableCascade combined: {sd_model.__class__.__name__}') - return sd_model + def cascade_post_load(sd_model): sd_model.prior_pipe.scheduler.config.clip_sample = False sd_model.default_scheduler = copy.deepcopy(sd_model.prior_pipe.scheduler) @@ -160,9 +162,9 @@ def cascade_post_load(sd_model): text_encoder=None, latent_dim_scale=sd_model.decoder_pipe.config.latent_dim_scale, ) - return sd_model + # Custom sampler support. Remove after the changes gets upstreamed: https://github.com/huggingface/diffusers/pull/9132 class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): def guidance_scale(self): diff --git a/modules/shared.py b/modules/shared.py index 544d3a07d..78d948cbc 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -615,8 +615,8 @@ options_templates.update(options_section(('saving-images', "Image Options"), { "image_sep_metadata": OptionInfo("