[Modular]z-image

[yiyixuxu]

import torch
from diffusers import ModularPipeline


# 1. Load the pipeline
# Use bfloat16 for optimal performance on supported GPUs
pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo")
print(f" pipe: {pipe}")
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")
# [Optional] Attention Backend
# Diffusers uses SDPA by default. Switch to Flash Attention for better efficiency if supported:
# pipe.transformer.set_attention_backend("flash")    # Enable Flash-Attention-2
# pipe.transformer.set_attention_backend("_flash_3") # Enable Flash-Attention-3

# [Optional] Model Compilation
# Compiling the DiT model accelerates inference, but the first run will take longer to compile.
# pipe.transformer.compile()

# [Optional] CPU Offloading
# Enable CPU offloading for memory-constrained devices.
# pipe.enable_model_cpu_offload()

prompt = "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."

# 2. Generate Image
## text-to-image
out1 = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

out1.save("t2i_output.png")

## image-to-image
prompt = "Young Chinese woman in red Hanfu, holding an apple"
strength = 0.8
out2 = pipe(
    image=out1,
    prompt=prompt,
    strength=strength,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]
out2.save(f"i2i_output_strength_{strength}.png")

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[asomoza]

here the previous models had a denoise in the name of the blocks, this is useful for adding the preview block with Mellon

[asomoza]

actually after looking more into it, all the other models had just one CoreDenoiseStep but this one has a separate one for t2i and i2i, so that's also a problem. I'll skip the preview block for now

[yiyixuxu]

yeah we will make it work, it all has a denoise block, just have to do a search until we find it. it is better if we goive more flexibility in terms on how one should pack the blocks, currently, the only constraint is to have top-level auto blocks like this

AUTO_BLOCKS = InsertableDict(
    [
        ("text_encoder", ...),
        ("vae_image_encoder", ...),
        ("denoise", ...),
        ("decode", ...),
    ]
)