Rename params for flux and flux vae, add comments explaining use of the config_path in model config

Author: brandonrising

invokeai/backend/flux/util.py

@@ -25,82 +25,48 @@ class ModelSpec:
 }
 
 
-ae_params = AutoEncoderParams(
-    resolution=256,
-    in_channels=3,
-    ch=128,
-    out_ch=3,
-    ch_mult=[1, 2, 4, 4],
-    num_res_blocks=2,
-    z_channels=16,
-    scale_factor=0.3611,
-    shift_factor=0.1159,
-)
+ae_params = {
+    "flux": AutoEncoderParams(
+        resolution=256,
+        in_channels=3,
+        ch=128,
+        out_ch=3,
+        ch_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        z_channels=16,
+        scale_factor=0.3611,
+        shift_factor=0.1159,
+    )
+}
 
 
-configs = {
-    "flux-dev": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-dev",
-        repo_flow="flux1-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV"),
-        params=FluxParams(
-            in_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
+params = {
+    "flux-dev": FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
     ),
-    "flux-schnell": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-schnell",
-        repo_flow="flux1-schnell.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_SCHNELL"),
-        params=FluxParams(
-            in_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=False,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
+    "flux-schnell": FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
     ),
 }

invokeai/backend/model_manager/load/model_loaders/flux.py

@@ -12,7 +12,7 @@
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.flux.model import Flux
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder
-from invokeai.backend.flux.util import ae_params, configs
+from invokeai.backend.flux.util import ae_params, params
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
@@ -59,7 +59,7 @@ def _load_model(
         model_path = Path(config.path)
 
         with SilenceWarnings():
-            model = AutoEncoder(ae_params)
+            model = AutoEncoder(ae_params[config.config_path])
             sd = load_file(model_path)
             model.load_state_dict(sd, assign=True)
             model.to(dtype=self._torch_dtype)
@@ -188,7 +188,7 @@ def _load_from_singlefile(
         model_path = Path(config.path)
 
         with SilenceWarnings():
-            model = Flux(configs[config.config_path].params)
+            model = Flux(params[config.config_path])
             sd = load_file(model_path)
             model.load_state_dict(sd, assign=True)
         return model
@@ -227,7 +227,7 @@ def _load_from_singlefile(
 
         with SilenceWarnings():
             with accelerate.init_empty_weights():
-                model = Flux(configs[config.config_path].params)
+                model = Flux(params[config.config_path])
                 model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=torch.bfloat16)
             sd = load_file(model_path)
             model.load_state_dict(sd, assign=True)

invokeai/backend/model_manager/probe.py

@@ -329,8 +329,16 @@ def _get_checkpoint_config_path(
                 checkpoint = ModelProbe._scan_and_load_checkpoint(model_path)
                 state_dict = checkpoint.get("state_dict") or checkpoint
                 if "guidance_in.out_layer.weight" in state_dict:
+                    # For flux, this is a key in invokeai.backend.flux.util.params
+                    #   Due to model type and format being the descriminator for model configs this
+                    #   is used rather than attempting to support flux with separate model types and format
+                    #   If changed in the future, please fix me
                     config_file = "flux-dev"
                 else:
+                    # For flux, this is a key in invokeai.backend.flux.util.params
+                    #   Due to model type and format being the descriminator for model configs this
+                    #   is used rather than attempting to support flux with separate model types and format
+                    #   If changed in the future, please fix me
                     config_file = "flux-schnell"
             else:
                 config_file = LEGACY_CONFIGS[base_type][variant_type]
@@ -345,7 +353,11 @@ def _get_checkpoint_config_path(
             )
         elif model_type is ModelType.VAE:
             config_file = (
-                "flux/flux1-vae.yaml"
+                # For flux, this is a key in invokeai.backend.flux.util.ae_params
+                #   Due to model type and format being the descriminator for model configs this
+                #   is used rather than attempting to support flux with separate model types and format
+                #   If changed in the future, please fix me
+                "flux"
                 if base_type is BaseModelType.Flux
                 else "stable-diffusion/v1-inference.yaml"
                 if base_type is BaseModelType.StableDiffusion1

invokeai/backend/quantization/scripts/load_flux_model_bnb_llm_int8.py

@@ -4,7 +4,7 @@
 from safetensors.torch import load_file, save_file
 
 from invokeai.backend.flux.model import Flux
-from invokeai.backend.flux.util import configs as flux_configs
+from invokeai.backend.flux.util import params
 from invokeai.backend.quantization.bnb_llm_int8 import quantize_model_llm_int8
 from invokeai.backend.quantization.scripts.load_flux_model_bnb_nf4 import log_time
 
@@ -22,11 +22,11 @@ def main():
 
     with log_time("Intialize FLUX transformer on meta device"):
         # TODO(ryand): Determine if this is a schnell model or a dev model and load the appropriate config.
-        params = flux_configs["flux-schnell"].params
+        p = params["flux-schnell"]
 
         # Initialize the model on the "meta" device.
         with accelerate.init_empty_weights():
-            model = Flux(params)
+            model = Flux(p)
 
     # TODO(ryand): We may want to add some modules to not quantize here (e.g. the proj_out layer). See the accelerate
     # `get_keys_to_not_convert(...)` function for a heuristic to determine which modules to not quantize.

invokeai/backend/quantization/scripts/load_flux_model_bnb_nf4.py

@@ -7,7 +7,7 @@
 from safetensors.torch import load_file, save_file
 
 from invokeai.backend.flux.model import Flux
-from invokeai.backend.flux.util import configs as flux_configs
+from invokeai.backend.flux.util import params
 from invokeai.backend.quantization.bnb_nf4 import quantize_model_nf4
 
 
@@ -35,11 +35,11 @@ def main():
     # inference_dtype = torch.bfloat16
     with log_time("Intialize FLUX transformer on meta device"):
         # TODO(ryand): Determine if this is a schnell model or a dev model and load the appropriate config.
-        params = flux_configs["flux-schnell"].params
+        p = params["flux-schnell"]
 
         # Initialize the model on the "meta" device.
         with accelerate.init_empty_weights():
-            model = Flux(params)
+            model = Flux(p)
 
     # TODO(ryand): We may want to add some modules to not quantize here (e.g. the proj_out layer). See the accelerate
     # `get_keys_to_not_convert(...)` function for a heuristic to determine which modules to not quantize.