# YiYi Notes: use this to generate mellon config
# TODO: 
# 1. the whole process to create a custom node should be automated with something like `node.push_to_hub()`
# 2. should be able to edit mellon_config.json


from diffusers.modular_pipelines.node_utils import ModularNode
from diffusers.modular_pipelines import ModularPipelineMixin
diffdiff = ModularPipelineMixin.from_pretrained("YiYiXu/modular-diffdiff", trust_remote_code=True)

# this is the complete pipeline
node = ModularNode(diffdiff)
node.save_mellon_config("diffdiff_mellon_config.json")


# this is the denoise pipeline, to be used in the modular setting, 
# remove the text_encoder, decode, and ip_adapter blocks since we already have these nodes in official modular diffusers mellon
diffdiff.blocks.pop("text_encoder")
diffdiff.blocks.pop("decode")
diffdiff.blocks.pop("ip_adapter")

node = ModularNode(diffdiff, label = "DiffDiffDenoiseBlocks")
node.save_mellon_config("diffdiff_denoise_mellon_config.json")