flux qint auto-download quantization map

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 = {}

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

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

modules/shared.py

@@ -615,8 +615,8 @@ options_templates.update(options_section(('saving-images', "Image Options"), {
 
     "image_sep_metadata": OptionInfo("<h2>Metadata/Logging</h2>", "", gr.HTML),
     "image_metadata": OptionInfo(True, "Include metadata"),
-    "save_txt": OptionInfo(False, "Create info file per image"),
-    "save_log_fn": OptionInfo("", "Update JSON log file per image", component_args=hide_dirs),
+    "save_txt": OptionInfo(False, "Create image info text file"),
+    "save_log_fn": OptionInfo("", "Append image info JSON file", component_args=hide_dirs),
     "image_sep_grid": OptionInfo("<h2>Grid Options</h2>", "", gr.HTML),
     "grid_save": OptionInfo(True, "Save all generated image grids"),
     "grid_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2"]}),