diff-diff.py

import logging
from mellon.NodeBase import NodeBase

# Configure logger
logger = logging.getLogger("mellon")
logger.setLevel(logging.DEBUG)




# Initialize components manager
# components = ComponentsManager()
from custom import components


class DiffDiff(NodeBase):
    def __init__(self, node_id=None):
        super().__init__(node_id)
        from diffusers.modular_pipelines.node_utils import ModularNode
        self._diffdiff_block = ModularNode.from_pretrained("YiYiXu/modular-diffdiff", trust_remote_code=True)
        self._diffdiff_block.setup(components=components)

    def execute(self, **kwargs):
        return self._diffdiff_block.execute(**kwargs)


class DiffDiffDenoise(NodeBase):
    def __init__(self, node_id=None):
        super().__init__(node_id)
        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)
        diffdiff.blocks.pop("text_encoder")
        diffdiff.blocks.pop("decode")
        diffdiff.blocks.pop("ip_adapter")
        self._diffdiff_block = ModularNode(diffdiff)
        self._diffdiff_block.setup(components=components)

    def execute(self, **kwargs):
        return self._diffdiff_block.execute(**kwargs)