YiYiXu/image_z_modular

how the mellon config is saved

from diffusers.modular_pipelines.mellon_node_utils import MellonParam, MellonPipelineConfig
ZIMAGE_NODE_SPECS = {
    "controlnet": None,
    "denoise": {
        "inputs": [
            MellonParam.embeddings(display="input"),
            MellonParam.width(),
            MellonParam.height(),
            MellonParam.seed(),
            MellonParam.num_inference_steps(default=9),
            MellonParam.image_latents_with_strength(),
            MellonParam.strength(),
        ],
        "model_inputs": [
            MellonParam.unet(),
            MellonParam.scheduler(),
        ],
        "outputs": [
            MellonParam.latents(display="output"),
            MellonParam.doc(),
        ],
        "required_inputs": ["embeddings"],
        "required_model_inputs": ["unet", "scheduler"],
        "block_name": "denoise",
    },
    "vae_encoder": {
        "inputs": [
            MellonParam.image(),
        ],
        "model_inputs": [
            MellonParam.vae(),
        ],
        "outputs": [
            MellonParam.image_latents(display="output"),
            MellonParam.doc(),
        ],
        "required_inputs": ["image"],
        "required_model_inputs": ["vae"],
        "block_name": "vae_encoder",
    },
    "text_encoder": {
        "inputs": [
            MellonParam.prompt(),
        ],
        "model_inputs": [
            MellonParam.text_encoders(),
        ],
        "outputs": [
            MellonParam.embeddings(display="output"),
            MellonParam.doc(),
        ],
        "required_inputs": ["prompt"],
        "required_model_inputs": ["text_encoders"],
        "block_name": "text_encoder",
    },
    "decoder": {
        "inputs": [
            MellonParam.latents(display="input"),
        ],
        "model_inputs": [
            MellonParam.vae(),
        ],
        "outputs": [
            MellonParam.images(),
            MellonParam.doc(),
        ],
        "required_inputs": ["latents"],
        "required_model_inputs": ["vae"],
        "block_name": "decode",
    },
}

ZIMAGE_PIPELINE_CONFIG = MellonPipelineConfig(
    node_specs=ZIMAGE_NODE_SPECS,
    label="ZImage",
    default_repo="Tongyi-MAI/Z-Image-Turbo",
    default_dtype="bfloat16",
)

ZIMAGE_PIPELINE_CONFIG.save("YiYiXu/image_z_modular", push_to_hub=True)