53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
from diffusers import (
|
|
QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel,
|
|
)
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
class QwenImageTransformer2DModel(
|
|
DiffusersQwenImageTransformer2DModel, OstrisModelMixin
|
|
):
|
|
aitk_subfolder = "transformer"
|
|
aitk_config_repo = "Qwen/Qwen-Image"
|
|
|
|
aitk_comfy_repo = "Comfy-Org/Qwen-Image_ComfyUI"
|
|
# the comfy files use the diffusers key layout directly (no conversion);
|
|
# fp8mixed carries float8_e4m3fn comfy_quant markers that attach straight
|
|
# onto this class's modules
|
|
aitk_comfy_weight_names = {
|
|
"Qwen/Qwen-Image": [
|
|
"split_files/diffusion_models/qwen_image_fp8mixed.safetensors",
|
|
# raw fp8 cast (no markers, diffusers keys) — loads via from_single_file
|
|
"split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors",
|
|
"split_files/diffusion_models/qwen_image_bf16.safetensors",
|
|
],
|
|
}
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["transformer_blocks"]
|
|
|
|
@classmethod
|
|
def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None):
|
|
from safetensors import safe_open
|
|
|
|
with safe_open(file_path, framework="pt") as f:
|
|
has_markers = any(k.endswith(".comfy_quant") for k in f.keys())
|
|
if has_markers:
|
|
# comfy prequantized checkpoint (diffusers key layout): the mixin
|
|
# path attaches the quantized layers
|
|
return super()._load_single_file(
|
|
file_path, dtype, config_path=config_path, subfolder=subfolder
|
|
)
|
|
# other single-file checkpoints carry diffusers or original key
|
|
# layouts; diffusers' single-file machinery owns that conversion
|
|
model = cls.from_single_file(
|
|
file_path,
|
|
config=config_path if config_path is not None else cls.aitk_config_repo,
|
|
subfolder="transformer",
|
|
torch_dtype=dtype,
|
|
)
|
|
model.to(dtype)
|
|
return model
|