Add basic support for Z-Image Turbo txt2img

configs/lcm-models.txt

@@ -3,4 +3,5 @@ rupeshs/sdxs-512-0.9-orig-vae
 rupeshs/hyper-sd-sdxl-1-step
 rupeshs/SDXL-Lightning-2steps
 stabilityai/sdxl-turbo
-SimianLuo/LCM_Dreamshaper_v7
+SimianLuo/LCM_Dreamshaper_v7
+Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32

src/backend/lcm_text_to_image.py

@@ -30,7 +30,7 @@
 )
 from backend.pipelines.lcm_lora import get_lcm_lora_pipeline
 from constants import DEVICE, GGUF_THREADS
-from diffusers import LCMScheduler
+from diffusers import LCMScheduler, ZImagePipeline
 from image_ops import resize_pil_image
 from backend.openvino.ov_hc_stablediffusion_pipeline import OvHcLatentConsistency
 from backend.gguf.gguf_diffusion import (
@@ -440,8 +440,11 @@ def init(
                 if lcm_diffusion_setting.lora.enabled:
                     print("Warning: Lora models not supported on OpenVINO mode")
             elif not lcm_diffusion_setting.use_gguf_model:
-                adapters = self.pipeline.get_active_adapters()
-                print(f"Active adapters : {adapters}")
+                try:
+                    adapters = self.pipeline.get_active_adapters()
+                    print(f"Active adapters : {adapters}")
+                except:
+                    print(f"No active adapters on selected pipeline")
 
     def _get_timesteps(self):
         time_steps = self.pipeline.scheduler.config.get("timesteps")
@@ -604,6 +607,8 @@ def generate(
                 == DiffusionTask.text_to_image.value
             ):
                 print(f"Using {self.pipeline.__class__.__name__}")
+                if self.pipeline.__class__.__name__ != "ZImagePipeline":
+                    pipeline_extra_args["timesteps"] = self._get_timesteps()
                 result_images = self.pipeline(
                     prompt=lcm_diffusion_setting.prompt,
                     negative_prompt=lcm_diffusion_setting.negative_prompt,
@@ -612,7 +617,7 @@ def generate(
                     width=lcm_diffusion_setting.image_width,
                     height=lcm_diffusion_setting.image_height,
                     num_images_per_prompt=lcm_diffusion_setting.number_of_images,
-                    timesteps=self._get_timesteps(),
+                    # timesteps=self._get_timesteps(),
                     **pipeline_extra_args,
                     **controlnet_args,
                 ).images

src/backend/pipelines/lcm.py

@@ -17,8 +17,10 @@
     AutoPipelineForImage2Image,
     StableDiffusionControlNetPipeline,
     StableDiffusionXLControlNetPipeline,
+    ZImagePipeline,
 )
 import pathlib
+from sdnq import SDNQConfig
 
 
 def _get_lcm_pipeline_from_base_model(
@@ -49,6 +51,9 @@ def load_taesd(
     use_local_model: bool = False,
     torch_data_type: torch.dtype = torch.float32,
 ):
+    pipeline_class = pipeline.__class__.__name__
+    if pipeline_class == "ZImagePipeline":
+        return
     tiny_vae = get_tiny_autoencoder_repo_id(pipeline.__class__.__name__)
     pipeline.vae = AutoencoderTiny.from_pretrained(
         tiny_vae,
@@ -76,6 +81,12 @@ def get_lcm_model_pipeline(
             "segmind/SSD-1B",
             use_local_model,
         )
+    elif "Z-Image" in model_id:
+        pipeline = ZImagePipeline.from_pretrained(
+            model_id,
+            local_files_only=use_local_model,
+            **pipeline_args,
+        )
     elif pathlib.Path(model_id).suffix == ".safetensors":
         # When loading a .safetensors model, the pipeline has to be created
         # with StableDiffusionPipeline() since it's the only class that
@@ -119,5 +130,7 @@ def get_image_to_image_pipeline(pipeline: Any) -> Any:
         return StableDiffusionImg2ImgPipeline(**components)
     elif pipeline_class == "StableDiffusionXLPipeline":
         return StableDiffusionXLImg2ImgPipeline(**components)
+    elif pipeline_class == "ZImagePipeline":
+        return pipeline
     else:
         raise Exception(f"Unknown pipeline {pipeline_class}")