Merge branch 'develop' into add_audio_to_output

Author: hadxp

diffusers_helper/lora_utils_kohya_ss/enums.py

@@ -1,85 +1,4 @@
-import sys
+from shared import LoraLoader
 
-if sys.version_info >= (3, 11):
-    from enum import StrEnum
-    # StrEnum is introduced in 3.11 while we support python 3.10
-else:
-    from enum import Enum, auto
-    from typing import Any
-
-    # Fallback for Python 3.10 and earlier
-    class StrEnum(str, Enum):
-        def __new__(cls, value, *args, **kwargs):
-            if not isinstance(value, (str, auto)):
-                raise TypeError(
-                    f"Values of StrEnums must be strings: {value!r} is a {type(value)}"
-                )
-            return super().__new__(cls, value, *args, **kwargs)
-
-        def __str__(self):
-            return str(self.value)
-
-        @staticmethod
-        def _generate_next_value_(
-            name: str, start: int, count: int, last_values: list[Any]
-        ) -> str:
-            return name
-
-
-class LoraLoader(StrEnum):
-    DIFFUSERS = "diffusers"
-    LORA_READY = "lora_ready"
-    DEFAULT = LORA_READY
-
-    @staticmethod
-    def supported_values() -> list[str]:
-        """Returns a list of all supported LoraLoader values."""
-        return [loader.value for loader in LoraLoader]
-
-    @staticmethod
-    def safe_parse(value: "str | LoraLoader") -> "LoraLoader":
-        if isinstance(value, LoraLoader):
-            return value
-        try:
-            return LoraLoader(value)
-        except ValueError:
-            return LoraLoader.DEFAULT
-
-
-if __name__ == "__main__":
-    # Test the StrEnum functionality
-    print("diffusers:", LoraLoader.DIFFUSERS)  # Should print "diffusers"
-    print("lora_ready:", LoraLoader.LORA_READY)  # Should print "lora_ready"
-    print("default:", LoraLoader.DEFAULT)  # Should print "lora_ready"
-    print(  # Should print all unique supported values (excludes aliases like DEFAULT)
-        "supported_values:", LoraLoader.supported_values()
-    )
-    try:
-        print("fail:", LoraLoader("invalid"))  # Should raise ValueError
-    except ValueError as e:
-        print("pass:", e)  # Prints: Invalid LoraLoader value: invalid
-    try:
-        print("pass:", LoraLoader("diffusers"))  # Should return LoraLoader.DIFFUSERS
-    except ValueError as e:
-        print("fail:", e)
-    try:
-        print("type of LoraLoader.DEFAULT:", type(LoraLoader.DEFAULT))
-        default = LoraLoader.DEFAULT
-        print("type of default:", type(default))  # Should be LoraLoader, not str
-    except Exception as e:
-        print(f"fail: {e}")
-
-    assert isinstance(LoraLoader("lora_ready"), StrEnum)
-    assert isinstance(LoraLoader.DIFFUSERS, LoraLoader), (
-        "DIFFUSERS should be an instance of LoraLoader"
-    )
-    assert LoraLoader.DEFAULT == LoraLoader.DIFFUSERS, (
-        "Default loader should be DIFFUSERS"
-    )
-    assert LoraLoader.DIFFUSERS != LoraLoader.LORA_READY, (
-        "DIFFUSERS should not equal LORA_READY"
-    )
-
-    assert LoraLoader.LORA_READY.value == "lora_ready", (
-        "lora_ready string should equal LoraLoader.LORA_READY"
-    )
+# todo: remove this import when the diffusers_helper is updated to use the new enums directly
+__all__ = ["LoraLoader"]

modules/generators/base_generator.py

@@ -5,6 +5,7 @@
 from diffusers_helper import lora_utils
 from typing import List, Optional, cast
 from pathlib import Path
+from transformers import BitsAndBytesConfig
 
 from diffusers_helper.lora_utils_kohya_ss.enums import LoraLoader
 from diffusers_helper.models.hunyuan_video_packed import (
@@ -13,6 +14,7 @@
 
 from ..settings import Settings
 from .model_configuration import ModelConfiguration
+from shared import QuantizationFormat
 
 # cSpell: ignore loras
 
@@ -23,6 +25,9 @@ class BaseModelGenerator(ABC):
     This defines the common interface that all model generators must implement.
     """
 
+    quantization_format: QuantizationFormat = QuantizationFormat.DEFAULT
+    quantization_config: BitsAndBytesConfig | None = None
+
     def __init__(
         self,
         text_encoder,
@@ -72,9 +77,32 @@ def __init__(
         self.gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.cpu = torch.device("cpu")
 
+        # quantization is currently global, configured in settings
+        # maybe add kwargs if we need this to be more dynamic per job?
+        self.quantization_format = self.settings.get(
+            "quantization_format", QuantizationFormat.integer_8bit
+        )
+        self.set_quantization_config()
+
         self.previous_model_hash: str = ""
         self.previous_model_configuration: ModelConfiguration | None = None
 
+    def set_quantization_config(self):
+        if self.quantization_format == QuantizationFormat.brain_floating_point_16bit:
+            # BF16 does not require a special config
+            pass
+        if self.quantization_format == QuantizationFormat.normal_float_4bit:
+            # 4-bit NF4 quantization config
+            self.quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype="bfloat16",
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+            )
+        if self.quantization_format == QuantizationFormat.integer_8bit:
+            # 8-bit integer quantization config
+            self.quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
     @abstractmethod
     def load_model(self) -> HunyuanVideoTransformer3DModelPacked:
         """
@@ -389,6 +417,7 @@ def load_loras(
 
         active_model_configuration = ModelConfiguration.from_lora_names_and_weights(
             self.get_model_name(),
+            self.quantization_format,
             selected_loras,
             selected_lora_values,
             self.settings.lora_loader,

modules/generators/f1_generator.py

@@ -4,6 +4,7 @@
 )
 from diffusers_helper.memory import DynamicSwapInstaller
 from .base_generator import BaseModelGenerator
+from shared import QuantizationFormat, timer
 
 
 class F1ModelGenerator(BaseModelGenerator):
@@ -28,6 +29,7 @@ def get_model_name(self):
         """
         return self.model_name
 
+    @timer
     def load_model(self):
         """
         Load the F1 transformer model.
@@ -44,12 +46,15 @@ def load_model(self):
 
         # Create the transformer model
         self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
-            path_to_load, torch_dtype=torch.bfloat16
+            path_to_load,
+            torch_dtype=torch.bfloat16,
+            quantization_config=self.quantization_config,
         ).cpu()
 
         # Configure the model
         self.transformer.eval()
-        self.transformer.to(dtype=torch.bfloat16)
+        if self.quantization_format == QuantizationFormat.brain_floating_point_16bit:
+            self.transformer.to(dtype=torch.bfloat16)
         self.transformer.requires_grad_(False)
 
         # Set up dynamic swap if not in high VRAM mode

modules/generators/model_configuration.py

@@ -304,6 +304,7 @@ def add_lora_setting(self, setting: ModelLoraSetting) -> None:
 @dataclass
 class ModelConfiguration:
     model_name: str
+    quantization_format: str
     settings: ModelSettings = field(default_factory=ModelSettings)
 
     @property
@@ -330,7 +331,9 @@ def validate(self) -> bool:
         return valid
 
     @staticmethod
-    def from_settings(model_name: str, settings: ModelSettings | dict | None):
+    def from_settings(
+        model_name: str, quantization_format: str, settings: ModelSettings | dict | None
+    ):
         model_settings: ModelSettings | None = None
         if settings is None:
             model_settings = ModelSettings()
@@ -344,11 +347,16 @@ def from_settings(model_name: str, settings: ModelSettings | dict | None):
         if model_settings is None:
             raise ValueError("Invalid config type for ModelConfiguration")
 
-        return ModelConfiguration(model_name=model_name, settings=model_settings)
+        return ModelConfiguration(
+            model_name=model_name,
+            quantization_format=quantization_format,
+            settings=model_settings,
+        )
 
     @staticmethod
     def from_lora_names_and_weights(
         model_name: str,
+        quantization_format: str,
         lora_names: list[str],
         lora_weights: list[float | int],
         lora_loader: str | LoraLoader,
@@ -374,7 +382,9 @@ def from_lora_names_and_weights(
             lora_settings=lora_settings, lora_loader=str(lora_loader)
         )
         return ModelConfiguration.from_settings(
-            model_name=model_name, settings=model_settings
+            model_name=model_name,
+            quantization_format=quantization_format,
+            settings=model_settings,
         )
 
     def set_model_name(self, model_name: str) -> "ModelConfiguration":

modules/generators/original_generator.py

@@ -4,6 +4,7 @@
 )
 from diffusers_helper.memory import DynamicSwapInstaller
 from .base_generator import BaseModelGenerator
+from shared import QuantizationFormat, timer
 
 
 class OriginalModelGenerator(BaseModelGenerator):
@@ -26,6 +27,7 @@ def get_model_name(self):
         """
         return self.model_name
 
+    @timer
     def load_model(self):
         """
         Load the Original transformer model.
@@ -42,12 +44,15 @@ def load_model(self):
 
         # Create the transformer model
         self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
-            path_to_load, torch_dtype=torch.bfloat16
+            path_to_load,
+            torch_dtype=torch.bfloat16,
+            quantization_config=self.quantization_config,
         ).cpu()
 
         # Configure the model
         self.transformer.eval()
-        self.transformer.to(dtype=torch.bfloat16)
+        if self.quantization_format == QuantizationFormat.brain_floating_point_16bit:
+            self.transformer.to(dtype=torch.bfloat16)
         self.transformer.requires_grad_(False)
 
         # Set up dynamic swap if not in high VRAM mode

modules/generators/video_base_generator.py

@@ -14,6 +14,7 @@
 from diffusers_helper.bucket_tools import find_nearest_bucket
 from diffusers_helper.hunyuan import vae_encode
 from .base_generator import BaseModelGenerator
+from shared import QuantizationFormat, timer
 
 
 class VideoBaseModelGenerator(BaseModelGenerator):
@@ -55,6 +56,7 @@ def get_model_name(self):
         """
         return self.model_name
 
+    @timer
     def load_model(self):
         """
         Load the Video transformer model.
@@ -71,12 +73,15 @@ def load_model(self):
 
         # Create the transformer model
         self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
-            path_to_load, torch_dtype=torch.bfloat16
+            path_to_load,
+            torch_dtype=torch.bfloat16,
+            quantization_config=self.quantization_config,
         ).cpu()
 
         # Configure the model
         self.transformer.eval()
-        self.transformer.to(dtype=torch.bfloat16)
+        if self.quantization_format == QuantizationFormat.brain_floating_point_16bit:
+            self.transformer.to(dtype=torch.bfloat16)
         self.transformer.requires_grad_(False)
 
         # Set up dynamic swap if not in high VRAM mode
@@ -585,8 +590,8 @@ def combine_videos(self, source_video_path, generated_video_path, output_path):
             )
 
             # Get the ffmpeg executable from the VideoProcessor class
-            from modules.toolbox.toolbox_processor import VideoProcessor
             from modules.toolbox.message_manager import MessageManager
+            from modules.toolbox.toolbox_processor import VideoProcessor
 
             # Create a message manager for logging
             message_manager = MessageManager()

modules/interface.py

@@ -334,13 +334,16 @@ def apply_startup_settings():
         )
         connect_audio_events(a, settings)
 
-        def refresh_loras():
+        def refresh_loras(current_selected):
             if enumerate_lora_dir_fn:
                 new_lora_names = enumerate_lora_dir_fn()
-                return gr.update(choices=new_lora_names)
+                preserved = [name for name in (current_selected or []) if name in new_lora_names]
+                return gr.update(choices=new_lora_names, value=preserved)
             return gr.update()
 
-        g["refresh_loras_button"].click(fn=refresh_loras, outputs=[g["lora_selector"]])
+        g["refresh_loras_button"].click(
+            fn=refresh_loras, inputs=[g["lora_selector"]], outputs=[g["lora_selector"]]
+        )
 
         # General Connections
         def initial_gallery_load():

modules/pipelines/worker.py

@@ -33,6 +33,7 @@
 from modules.llm_enhancer import unload_enhancing_model
 from . import create_pipeline
 from modules.studio_manager import StudioManager
+from shared import timer
 
 # cSpell: disable hunyan, loras
 
@@ -87,6 +88,7 @@ def get_cached_or_encode_prompt(
 
 
 @torch.no_grad()
+@timer
 def worker(
     model_type,
     input_image,

modules/settings.py

@@ -2,7 +2,8 @@
 from pathlib import Path
 from typing import Dict, Any, Optional
 import os
-from diffusers_helper.lora_utils_kohya_ss.enums import LoraLoader
+
+from shared import LoraLoader, QuantizationFormat
 
 
 class Settings:
@@ -45,6 +46,7 @@ def __init__(self):
 Enhanced prompt:""",
             "lora_loader": LoraLoader.DEFAULT,  # lora_loader options: diffusers, lora_ready. DEFAULT is existing behavior of diffusers
             "reuse_model_instance": False,  # Reuse model instance across generations - default of False is existing behavior
+            "quantization_format": QuantizationFormat.DEFAULT,  # Default quantization format
         }
         self.settings = self.load_settings()
 
@@ -73,6 +75,21 @@ def lora_loader(self, value: str | LoraLoader):
     def reuse_model_instance(self) -> bool:
         return self.settings.get("reuse_model_instance", False)
 
+    @property
+    def quantization_format(self) -> QuantizationFormat:
+        return QuantizationFormat.safe_parse(
+            self.settings.get("quantization_format", QuantizationFormat.DEFAULT)
+        )
+
+    @quantization_format.setter
+    def quantization_format(self, value: str | QuantizationFormat):
+        if not value:
+            value = QuantizationFormat.DEFAULT
+        if isinstance(value, str):
+            value = QuantizationFormat.safe_parse(value)
+
+        self.set("quantization_format", value)
+
     def load_settings(self) -> Dict[str, Any]:
         """Load settings from file or return defaults"""
         if self.settings_file.exists():