github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Refactor qwen5b model code to be qwen 5b specific

Browse Source
This commit is contained in:
Jaret Burkett
2025-08-06 10:54:56 -06:00
parent 5d8922fca2
commit 14ccf2f3ce
4 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
extensions_built_in/diffusion_models/__init__.py
View File

@@ -3,7 +3,7 @@ from .hidream import HidreamModel, HidreamE1Model
from .f_light import FLiteModel from .f_light import FLiteModel
from .omnigen2 import OmniGen2Model from .omnigen2 import OmniGen2Model
from .flux_kontext import FluxKontextModel from .flux_kontext import FluxKontextModel
from .wan22 import Wan22Model from .wan22 import Wan225bModel
from .qwen_image import QwenImageModel from .qwen_image import QwenImageModel
AI_TOOLKIT_MODELS = [ AI_TOOLKIT_MODELS = [
@@ -14,6 +14,6 @@ AI_TOOLKIT_MODELS = [
FLiteModel, FLiteModel,
OmniGen2Model, OmniGen2Model,
FluxKontextModel, FluxKontextModel,
Wan22Model, Wan225bModel,
QwenImageModel, QwenImageModel,
] ]

2
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -136,7 +136,7 @@ class QwenImageModel(BaseModel):
if self.model_config.quantize_te: if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder") self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype( quantize(text_encoder, weights=get_qtype(
self.model_config.qtype)) self.model_config.qtype_te))
freeze(text_encoder) freeze(text_encoder)
flush() flush()

2
extensions_built_in/diffusion_models/wan22/__init__.py
View File

@@ -1 +1 @@
from .wan22_model import Wan22Model from .wan22_5b_model import Wan225bModel

2
extensions_built_in/diffusion_models/wan22/wan22_model.py → extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
View File

@@ -80,7 +80,7 @@ def time_text_monkeypatch(
return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
class Wan22Model(Wan21): class Wan225bModel(Wan21):
arch = "wan22_5b" arch = "wan22_5b"
_wan_generation_scheduler_config = scheduler_configUniPC _wan_generation_scheduler_config = scheduler_configUniPC
_wan_expand_timesteps = True _wan_expand_timesteps = True