support z-image with lora

Author: gushiqiao

lightx2v/models/runners/z_image/z_image_runner.py

@@ -7,6 +7,7 @@
 from loguru import logger
 
 from lightx2v.models.input_encoders.hf.z_image.qwen3_model import Qwen3Model_TextEncoder
+from lightx2v.models.networks.lora_adapter import LoraAdapter
 from lightx2v.models.networks.z_image.model import ZImageTransformerModel
 from lightx2v.models.runners.default_runner import DefaultRunner
 from lightx2v.models.schedulers.z_image.scheduler import ZImageScheduler
@@ -30,6 +31,24 @@ def calculate_dimensions(target_area, ratio):
     return width, height, None
 
 
+def build_z_image_model_with_lora(z_image_module, config, model_kwargs, lora_configs):
+    lora_dynamic_apply = config.get("lora_dynamic_apply", False)
+
+    if lora_dynamic_apply:
+        lora_path = lora_configs[0]["path"]
+        lora_strength = lora_configs[0]["strength"]
+        model_kwargs["lora_path"] = lora_path
+        model_kwargs["lora_strength"] = lora_strength
+        model = z_image_module(**model_kwargs)
+    else:
+        assert not config.get("dit_quantized", False), "Online LoRA only for quantized models; merging LoRA is unsupported."
+        assert not config.get("lazy_load", False), "Lazy load mode does not support LoRA merging."
+        model = z_image_module(**model_kwargs)
+        lora_adapter = LoraAdapter(model)
+        lora_adapter.apply_lora(lora_configs)
+    return model
+
+
 @RUNNER_REGISTER("z_image")
 class ZImageRunner(DefaultRunner):
     model_cpu_offload_seq = "text_encoder->transformer->vae"
@@ -45,7 +64,16 @@ def load_model(self):
         self.vae = self.load_vae()
 
     def load_transformer(self):
-        model = ZImageTransformerModel(os.path.join(self.config["model_path"], "transformer"), self.config, self.init_device)
+        z_image_model_kwargs = {
+            "model_path": os.path.join(self.config["model_path"], "transformer"),
+            "config": self.config,
+            "device": self.init_device,
+        }
+        lora_configs = self.config.get("lora_configs")
+        if not lora_configs:
+            model = ZImageTransformerModel(**z_image_model_kwargs)
+        else:
+            model = build_z_image_model_with_lora(ZImageTransformerModel, self.config, z_image_model_kwargs, lora_configs)
         return model
 
     def load_text_encoder(self):