From 14ccf2f3cea98aa70436f0553f0a5788ca55e7f8 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 6 Aug 2025 10:54:56 -0600
Subject: [PATCH] Refactor qwen5b model code to be qwen 5b specific

---
 extensions_built_in/diffusion_models/__init__.py              | 4 ++--
 extensions_built_in/diffusion_models/qwen_image/qwen_image.py | 2 +-
 extensions_built_in/diffusion_models/wan22/__init__.py        | 2 +-
 .../wan22/{wan22_model.py => wan22_5b_model.py}               | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)
 rename extensions_built_in/diffusion_models/wan22/{wan22_model.py => wan22_5b_model.py} (99%)

diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py
index 035a3188..04ff0ff3 100644
--- a/extensions_built_in/diffusion_models/__init__.py
+++ b/extensions_built_in/diffusion_models/__init__.py
@@ -3,7 +3,7 @@ from .hidream import HidreamModel, HidreamE1Model
 from .f_light import FLiteModel
 from .omnigen2 import OmniGen2Model
 from .flux_kontext import FluxKontextModel
-from .wan22 import Wan22Model
+from .wan22 import Wan225bModel
 from .qwen_image import QwenImageModel
 
 AI_TOOLKIT_MODELS = [
@@ -14,6 +14,6 @@ AI_TOOLKIT_MODELS = [
     FLiteModel,
     OmniGen2Model,
     FluxKontextModel,
-    Wan22Model,
+    Wan225bModel,
     QwenImageModel,
 ]
diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py
index f834a8a3..f7607ccb 100644
--- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py
+++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py
@@ -136,7 +136,7 @@ class QwenImageModel(BaseModel):
         if self.model_config.quantize_te:
             self.print_and_status_update("Quantizing Text Encoder")
             quantize(text_encoder, weights=get_qtype(
-                self.model_config.qtype))
+                self.model_config.qtype_te))
             freeze(text_encoder)
             flush()
 
diff --git a/extensions_built_in/diffusion_models/wan22/__init__.py b/extensions_built_in/diffusion_models/wan22/__init__.py
index 765b1a18..993acbfb 100644
--- a/extensions_built_in/diffusion_models/wan22/__init__.py
+++ b/extensions_built_in/diffusion_models/wan22/__init__.py
@@ -1 +1 @@
-from .wan22_model import Wan22Model
\ No newline at end of file
+from .wan22_5b_model import Wan225bModel
\ No newline at end of file
diff --git a/extensions_built_in/diffusion_models/wan22/wan22_model.py b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
similarity index 99%
rename from extensions_built_in/diffusion_models/wan22/wan22_model.py
rename to extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
index 328ea807..9c446cc5 100644
--- a/extensions_built_in/diffusion_models/wan22/wan22_model.py
+++ b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
@@ -80,7 +80,7 @@ def time_text_monkeypatch(
 
     return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
 
-class Wan22Model(Wan21):
+class Wan225bModel(Wan21):
     arch = "wan22_5b"
     _wan_generation_scheduler_config = scheduler_configUniPC
     _wan_expand_timesteps = True