import torch
from diffusers import ZImagePipeline, ZImageTransformer2DModel

# Load the custom transformer
transformer = ZImageTransformer2DModel.from_pretrained(
    "linoyts/beyond-reality-z-image-diffusers",
    torch_dtype=torch.bfloat16
)

# Load the pipeline with custom transformer
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

# Generate an image
prompt = "A beautiful landscape with mountains and a lake, photorealistic, 8k"
image = pipe(
    prompt=prompt,
    num_inference_steps=8,
    guidance_scale=0.0,  # Z-Image-Turbo uses guidance_scale=0
    width=1024,
    height=1024,
).images[0]

image.save("output.png")

Original Model

Source: Nurburgring/BEYOND_REALITY_Z_IMAGE
Base Architecture: Tongyi-MAI/Z-Image-Turbo

Conversion Details

The model was converted from ComfyUI format to diffusers format with the following key transformations:

Removed model.diffusion_model. prefix from all keys
Renamed x_embedder to all_x_embedder.2-1
Renamed final_layer to all_final_layer.2-1
Split attention.qkv into attention.to_q, attention.to_k, attention.to_v
Renamed attention.out to attention.to_out.0
Renamed attention.q_norm to attention.norm_q
Renamed attention.k_norm to attention.norm_k