1from diffusers.modular_pipelines.mellon_node_utils import MellonParam, MellonPipelineConfig
2ZIMAGE_NODE_SPECS = {
3 "controlnet": None,
4 "denoise": {
5 "inputs": [
6 MellonParam.embeddings(display="input"),
7 MellonParam.width(),
8 MellonParam.height(),
9 MellonParam.seed(),
10 MellonParam.num_inference_steps(default=9),
11 MellonParam.image_latents_with_strength(),
12 MellonParam.strength(),
13 ],
14 "model_inputs": [
15 MellonParam.unet(),
16 MellonParam.scheduler(),
17 ],
18 "outputs": [
19 MellonParam.latents(display="output"),
20 MellonParam.doc(),
21 ],
22 "required_inputs": ["embeddings"],
23 "required_model_inputs": ["unet", "scheduler"],
24 "block_name": "denoise",
25 },
26 "vae_encoder": {
27 "inputs": [
28 MellonParam.image(),
29 ],
30 "model_inputs": [
31 MellonParam.vae(),
32 ],
33 "outputs": [
34 MellonParam.image_latents(display="output"),
35 MellonParam.doc(),
36 ],
37 "required_inputs": ["image"],
38 "required_model_inputs": ["vae"],
39 "block_name": "vae_encoder",
40 },
41 "text_encoder": {
42 "inputs": [
43 MellonParam.prompt(),
44 ],
45 "model_inputs": [
46 MellonParam.text_encoders(),
47 ],
48 "outputs": [
49 MellonParam.embeddings(display="output"),
50 MellonParam.doc(),
51 ],
52 "required_inputs": ["prompt"],
53 "required_model_inputs": ["text_encoders"],
54 "block_name": "text_encoder",
55 },
56 "decoder": {
57 "inputs": [
58 MellonParam.latents(display="input"),
59 ],
60 "model_inputs": [
61 MellonParam.vae(),
62 ],
63 "outputs": [
64 MellonParam.images(),
65 MellonParam.doc(),
66 ],
67 "required_inputs": ["latents"],
68 "required_model_inputs": ["vae"],
69 "block_name": "decode",
70 },
71}
72
73ZIMAGE_PIPELINE_CONFIG = MellonPipelineConfig(
74 node_specs=ZIMAGE_NODE_SPECS,
75 label="ZImage",
76 default_repo="Tongyi-MAI/Z-Image-Turbo",
77 default_dtype="bfloat16",
78)
79
80ZIMAGE_PIPELINE_CONFIG.save("YiYiXu/image_z_modular", push_to_hub=True)