models/z_image/z_image_main.py

import json
import os
import torch
from accelerate import init_empty_weights
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import logging
from mmgp import offload
from shared.utils import files_locator as fl
from transformers import AutoTokenizer, Qwen3ForCausalLM

from .autoencoder_kl import AutoencoderKL
from .pipeline_z_image import ZImagePipeline
from .z_image_transformer2d import ZImageTransformer2DModel


logger = logging.get_logger(__name__)


def conv_state_dict(sd: dict) -> dict:
    if "x_embedder.weight" not in sd and "model.diffusion_model.x_embedder.weight" not in sd:
        return sd

    inverse_replace = {
        "final_layer.": "all_final_layer.2-1.",
        "x_embedder.": "all_x_embedder.2-1.",
        ".attention.out.bias": ".attention.to_out.0.bias",
        ".attention.k_norm.weight": ".attention.norm_k.weight",
        ".attention.q_norm.weight": ".attention.norm_q.weight",
        ".attention.out.weight": ".attention.to_out.0.weight",
    }

    out_sd: dict[str, torch.Tensor] = {}

    for key, tensor in sd.items():
        key = key.replace("model.diffusion_model.", "")
        
        new_key = key
        for ori_sub, orig_sub in inverse_replace.items():
            new_key = new_key.replace(ori_sub, orig_sub)
        out_sd[new_key] = tensor

    return out_sd


_ZIMAGE_FUSED_SPLIT_MAP = {
    "attention.to_qkv": {"mapped_modules": ("attention.to_q", "attention.to_k", "attention.to_v")},
    "attention.qkv": {"mapped_modules": ("attention.to_q", "attention.to_k", "attention.to_v")},
    "feed_forward.net.0.proj": {"mapped_modules": ("feed_forward.w3", "feed_forward.w1")},
    "feed_forward.net.2": {"mapped_modules": ("feed_forward.w2",)},
}


class model_factory:
    def __init__(
        self,
        checkpoint_dir,
        model_filename=None,
        model_type=None,
        model_def=None,
        base_model_type=None,
        text_encoder_filename=None,
        quantizeTransformer=False,
        dtype=torch.bfloat16,
        VAE_dtype=torch.float32,
        mixed_precision_transformer=False,
        save_quantized=False,
        is_control=False,
        **kwargs,
    ):
        model_def = model_def or {}
        source =  model_def.get("source", None)
        module_source =  model_def.get("module_source", None)


        # model_filename can be a string or list of files (transformer + modules)
        transformer_filename = model_filename[0] if isinstance(model_filename, (list, tuple)) else model_filename
        if transformer_filename is None:
            raise ValueError("No transformer filename provided for Z-Image.")

        self.base_model_type = base_model_type
        self.is_control = is_control
        self.model_def = model_def

        default_transformer_config = os.path.join(os.path.dirname(os.path.abspath(__file__)), "configs", f"{base_model_type}.json")

        def preprocess_sd(state_dict):
            return conv_state_dict(state_dict)

        model_class = ZImageTransformer2DModel

        kwargs_light= { "writable_tensors": False, "preprocess_sd": preprocess_sd, "fused_split_map": _ZIMAGE_FUSED_SPLIT_MAP }
        # model_filename contains all files to load (transformer + modules merged by loader)
        import json
        import accelerate
        with open(default_transformer_config, "r") as f:
            config = json.load(f)
        config.pop("_class_name", None)
        config.pop("_diffusers_version", None)

        with accelerate.init_empty_weights():
            transformer = model_class(**config)

        if source is not None:
            offload.load_model_data(transformer, fl.locate_file(source), **kwargs_light)
        elif module_source is not None:
            offload.load_model_data(transformer, model_filename[:1] + [fl.locate_file(module_source)], **kwargs_light)
        else:
            offload.load_model_data(transformer, model_filename, **kwargs_light)

        from wgp import save_model
        from mmgp.safetensors2 import torch_load_file

        transformer.to(dtype)

        if module_source is not None:
            save_model(transformer, model_type, dtype, None, is_module=True, filter=list(torch_load_file(fl.locate_file(module_source))), module_source_no=1)

        if not source is None:
            save_model(transformer, model_type, dtype, None, submodel_no= 1)

        if save_quantized:
            from wgp import save_quantized_model
            save_quantized_model(transformer, model_type, transformer_filename, dtype, default_transformer_config)

        # Text encoder

        # text_encoder = Qwen3ForCausalLM.from_pretrained(os.path.dirname(text_encoder_filename), trust_remote_code=True)
        # text_encoder.to(torch.bfloat16)
        # offload.save_model(text_encoder, "c:/temp/qwnen3_bf16_.safetensors")
        
        text_encoder = offload.fast_load_transformers_model( text_encoder_filename, writable_tensors=True, modelClass=Qwen3ForCausalLM,)

        # Tokenizer
        text_encoder_folder = model_def.get("text_encoder_folder")
        if text_encoder_folder:
            tokenizer_path = os.path.dirname(fl.locate_file(os.path.join(text_encoder_folder, "tokenizer_config.json")))
        else:
            tokenizer_path = os.path.dirname(text_encoder_filename)
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)

        # VAE
        vae_filename = fl.locate_file("ZImageTurbo_VAE_bf16.safetensors")
        vae_config_path = fl.locate_file("ZImageTurbo_VAE_bf16_config.json") 

        vae = offload.fast_load_transformers_model(
            vae_filename,
            writable_tensors=True,
            modelClass=AutoencoderKL,
            defaultConfigPath=vae_config_path,
            default_dtype=VAE_dtype,
        )

        # Scheduler
        with open(fl.locate_file("ZImageTurbo_scheduler_config.json"), "r", encoding="utf-8") as f:
            scheduler_config = json.load(f)

        scheduler = FlowMatchEulerDiscreteScheduler(**scheduler_config)

        self.pipeline = ZImagePipeline(
            scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer
        )
        self.transformer = transformer
        self.text_encoder = text_encoder
        self.tokenizer = tokenizer
        self.vae = vae
        self.scheduler = scheduler

    def generate(
        self,
        seed: int | None = None,
        input_prompt: str = "",
        n_prompt: str | None = None,
        sampling_steps: int = 20,
        sample_solver: str = "default",
        width: int = 1024,
        height: int = 1024,
        guide_scale: float = 0.0,
        batch_size: int = 1,
        callback=None,
        max_sequence_length: int = 512,
        VAE_tile_size=None,
        cfg_normalization: bool = False,
        cfg_truncation: float = 1.0,
        input_frames=None,
        input_masks=None,
        context_scale: float = [0],
        input_ref_images = None,
        NAG_scale: float = 1.0,
        NAG_tau: float = 3.5,
        NAG_alpha: float = 0.5,
        loras_slists=None,
        **kwargs,
    ):
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
        if seed is None or seed < 0:
            generator.seed()
        else:
            generator.manual_seed(int(seed))

        if VAE_tile_size is not None and hasattr(self.vae, "use_tiling"):
            if isinstance(VAE_tile_size, int):
                tiling = VAE_tile_size > 0
                tile_size = max(VAE_tile_size, 0)
            else:
                tiling = bool(VAE_tile_size[0])
                tile_size = VAE_tile_size[1] if len(VAE_tile_size) > 1 else 0
            self.vae.use_tiling = tiling
            self.vae.tile_latent_min_height = tile_size
            self.vae.tile_latent_min_width = tile_size

        unified_solver = self.model_def.get("unified_solver", False)
        if unified_solver:
            sample_solver = "unified"
        elif not sample_solver:
            sample_solver = "default"

        if self.model_def.get("guidance_max_phases", 0) < 1:
            guide_scale = 0

        images = self.pipeline(
            prompt=input_prompt,
            negative_prompt=n_prompt,
            num_inference_steps=sampling_steps,
            sample_solver=sample_solver,
            guidance_scale=guide_scale,
            num_images_per_prompt=batch_size,
            generator=generator,
            height=height,
            width=width,
            max_sequence_length=max_sequence_length,
            callback_on_step_end=None,
            output_type="pt",
            return_dict=True,
            cfg_normalization=cfg_normalization,
            cfg_truncation=cfg_truncation,
            callback=callback,
            pipeline=self.pipeline,
            control_image=input_frames,
            inpaint_mask=input_masks,
            control_context_scale=None if context_scale is None else context_scale[0],
            input_ref_images= input_ref_images,
            NAG_scale=NAG_scale,
            NAG_tau=NAG_tau,
            NAG_alpha=NAG_alpha,
            loras_slists=loras_slists,
        )

        if images is None:
            return None

        if not torch.is_tensor(images):
            images = torch.tensor(images)

        return images.transpose(0, 1)

    def get_loras_transformer(self, *args, **kwargs):
        return [], []

    @property
    def _interrupt(self):
        return getattr(self.pipeline, "_interrupt", False)

    @_interrupt.setter
    def _interrupt(self, value):
        if hasattr(self, "pipeline"):
            self.pipeline._interrupt = value