flux nf4 offline load

CHANGELOG.md

@@ -7,7 +7,7 @@
 Summer break is over and we are back with a massive update!  
 
 Support for all of the new models:  
-- [Black Forest Labs FLUX.1](https://blackforestlabs.ai/announcing-black-forest-labs/) original variations and multiple quantized variations (*qint8, qint4, nf4*)!  
+- [Black Forest Labs FLUX.1](https://blackforestlabs.ai/announcing-black-forest-labs/)  
 - [AuraFlow](https://huggingface.co/fal/AuraFlow)  
 - [AlphaVLLM Lumina-Next-SFT](https://huggingface.co/Alpha-VLLM/Lumina-Next-SFT-diffusers)  
 - [Kwai Kolors](https://huggingface.co/Kwai-Kolors/Kolors)  
@@ -15,7 +15,7 @@ Support for all of the new models:
 
 What else? Just a bit... ;)  
 
-New **fast-install** mode, new **ControlNet-Union** *all-in-one* model, support for **DoRA** networks, additional **VLM** models, new **AuraSR** upscaler, new **Optimum Quanto** and **BitsAndBytes** quantization modes, new **balanced offload** mode and more...  
+New **fast-install** mode, new **Optimum Quanto** and **BitsAndBytes** quantization modes, new **balanced offload** mode, new **ControlNet-Union** *all-in-one* model, support for **DoRA** networks, additional **VLM** models, new **AuraSR** upscaler,  and more...  
 
 **Breaking Changes...**
 
@@ -31,8 +31,7 @@ To use and of the new models, simply select model from *Networks -> Reference* a
 - [Black Forest Labs FLUX.1](https://blackforestlabs.ai/announcing-black-forest-labs/)  
   FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching  
   This is a very large model at ~32GB in size, its recommended to use a) offloading, b) quantization  
-  *Note*: [FLUX.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) variant is a gated model, you need to accept the terms and conditions to use it  
-  For more information, see [Wiki](https://github.com/vladmandic/automatic/wiki/FLUX)  
+  For more information on how to donwload and use FLUX.1, see [Wiki](https://github.com/vladmandic/automatic/wiki/FLUX)  
   SD.Next supports:  
   - [FLUX.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) and [FLUX.1 Schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) original variations  
   - additional [qint8](https://huggingface.co/Disty0/FLUX.1-dev-qint8) and [qint4](https://huggingface.co/Disty0/FLUX.1-dev-qint4) quantized variations  
@@ -59,6 +58,7 @@ To use and of the new models, simply select model from *Networks -> Reference* a
   balanced offload will dynamically split and offload models from the GPU based on the max configured GPU and CPU memory size  
   model parts that dont fit in the GPU will be dynamically sliced and offloaded to the CPU  
   see *Settings -> Diffusers Settings -> Max GPU memory and Max CPU memory*  
+  *note*: recommended value for max GPU memory is ~80% of your total GPU memory  
   *note*: balanced offload will force loading LoRA with Diffusers method  
   *note*: balanced offload is not compatible with Optimum Quanto  
 - support for **Optimum Quanto** with 8 bit and 4 bit quantization options, thanks @Disty0 and @Trojaner!  

extensions-builtin/Lora/extra_networks_lora.py

@@ -110,9 +110,10 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
                     if shared.opts.lora_fuse_diffusers:
                         shared.sd_model.unfuse_lora()
                     shared.sd_model.unload_lora_weights() # fails for non-CLIP models
-                    shared.log.debug("LoRA unload")
-                except Exception as e:
-                    shared.log.warning(f"LoRA unload: {e}")
+                    # shared.log.debug("LoRA unload")
+                except Exception:
+                    # shared.log.warning(f"LoRA unload: {e}")
+                    pass
         if not self.active and getattr(networks, "originals", None ) is not None:
             networks.originals.undo() # remove patches
             if networks.debug:

modules/model_flux_nf4.py

@@ -149,13 +149,16 @@ def create_quantized_param(
 
 
 def load_flux_nf4(checkpoint_info, diffusers_load_config):
-    if os.path.exists(checkpoint_info.path) and os.path.isfile(checkpoint_info.path):
-        ckpt_path = checkpoint_info.path
+    repo_path = checkpoint_info.path
+    if os.path.exists(repo_path) and os.path.isfile(repo_path):
+        ckpt_path = repo_path
+    if os.path.exists(repo_path) and os.path.isdir(repo_path) and os.path.exists(os.path.join(repo_path, "diffusion_pytorch_model.safetensors")):
+        ckpt_path = os.path.join(repo_path, "diffusion_pytorch_model.safetensors")
     else:
-        ckpt_path = hf_hub_download(checkpoint_info.path, filename="diffusion_pytorch_model.safetensors", cache_dir=shared.opts.diffusers_dir)
+        ckpt_path = hf_hub_download(repo_path, filename="diffusion_pytorch_model.safetensors", cache_dir=shared.opts.diffusers_dir)
     original_state_dict = safetensors.torch.load_file(ckpt_path)
 
-    if 'sayakpaul/flux.1-dev-nf4' in checkpoint_info.path:
+    if 'sayakpaul' in checkpoint_info.path:
         converted_state_dict = original_state_dict # already converted
     else:
         try:

modules/modelloader.py

@@ -279,7 +279,7 @@ def load_diffusers_models(clear=True):
                     mtime = os.path.getmtime(commit)
                     info = os.path.join(commit, "model_info.json")
                     index = os.path.join(commit, "model_index.json")
-                    if not os.path.exists(index):
+                    if (not os.path.exists(index)) and (not os.path.exists(info)):
                         debug(f'Diffusers skip model no info: {name}')
                         continue
                     repo = { 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': commit, 'hash': snapshot, 'mtime': mtime, 'model_info': info, 'model_index': index }

modules/sd_vae_taesd.py

@@ -50,48 +50,28 @@ def Encoder(latent_channels=4):
     )
 
 def Decoder(latent_channels=4):
-    return nn.Sequential(
-        Clamp(), conv(latent_channels, 64), nn.ReLU(),
-        Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-        Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-        Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-        Block(64, 64), conv(64, 3),
-    )
-
-
-class TAESD2(nn.Module): # pylint: disable=abstract-method
-    latent_magnitude = 3
-    latent_shift = 0.5
-
-    def __init__(self, encoder_path="taesd_encoder.pth", decoder_path="taesd_decoder.pth", latent_channels=None):
-        """Initialize pretrained TAESD on the given device from the given checkpoints."""
-        super().__init__()
-        if latent_channels is None:
-            latent_channels = self.guess_latent_channels(str(encoder_path))
-        self.encoder = Encoder(latent_channels)
-        self.decoder = Decoder(latent_channels)
-        if encoder_path is not None:
-            self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu"))
-        if decoder_path is not None:
-            self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu"))
-
-    def guess_latent_channels(self, encoder_path):
-        """guess latent channel count based on encoder filename"""
-        if "taef1" in encoder_path:
-            return 16
-        if "taesd3" in encoder_path:
-            return 16
-        return 4
-
-    @staticmethod
-    def scale_latents(x):
-        """raw latents -> [0, 1]"""
-        return x.div(2 * TAESD.latent_magnitude).add(TAESD.latent_shift).clamp(0, 1)
-
-    @staticmethod
-    def unscale_latents(x):
-        """[0, 1] -> raw latents"""
-        return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
+    from modules import shared
+    if shared.opts.live_preview_taesd_layers == 1:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
+    elif shared.opts.live_preview_taesd_layers == 2:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
+    else:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
 
 
 class TAESD(nn.Module): # pylint: disable=abstract-method

modules/shared.py

@@ -709,6 +709,7 @@ options_templates.update(options_section(('live-preview', "Live Previews"), {
     "show_progress_type": OptionInfo("Approximate", "Live preview method", gr.Radio, {"choices": ["Simple", "Approximate", "TAESD", "Full VAE"]}),
     "live_preview_content": OptionInfo("Combined", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"], "visible": False}),
     "live_preview_refresh_period": OptionInfo(500, "Progress update period", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}),
+    "live_preview_taesd_layers": OptionInfo(3, "TAESD decode layers", gr.Slider, {"minimum": 1, "maximum": 3, "step": 1}),
     "logmonitor_show": OptionInfo(True, "Show log view"),
     "logmonitor_refresh_period": OptionInfo(5000, "Log view update period", gr.Slider, {"minimum": 0, "maximum": 30000, "step": 25}),
 }))