diff --git a/configs/lcm-models.txt b/configs/lcm-models.txt index 0162aa4..46b0b3f 100644 --- a/configs/lcm-models.txt +++ b/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 \ No newline at end of file +SimianLuo/LCM_Dreamshaper_v7 +Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32 diff --git a/src/backend/lcm_text_to_image.py b/src/backend/lcm_text_to_image.py index 799a7d5..82b3397 100644 --- a/src/backend/lcm_text_to_image.py +++ b/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 diff --git a/src/backend/pipelines/lcm.py b/src/backend/pipelines/lcm.py index faa431e..dd3c2e6 100644 --- a/src/backend/pipelines/lcm.py +++ b/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}")