qwen_image_scale_rope.py

from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
import torch

class QwenImageScaleRope(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="QwenImageScaleRope",
            category="_for_testing",
            inputs=[
                io.Model.Input("model"),
                io.Int.Input("scale", default=1, min=1, max=10, step=1),
            ],
            outputs=[
                io.Model.Output(),
            ],
            description="A node that adds a wrapper to scale the RoPE embeddings in Qwen image attention.",
        )

    @classmethod
    def execute(cls, model, scale: int) -> io.NodeOutput:
        new_model = model.clone()
        def patch(args, original_func):
            args["pe"] = torch.linalg.matrix_power(args["pe"], scale)
            return original_func["original_block"](args)
        new_model.set_model_patch_replace(patch, "dit", "double_block", 0)
        return io.NodeOutput(new_model)

class QwenImageScaleRopeExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            QwenImageScaleRope,
        ]

async def comfy_entrypoint() -> QwenImageScaleRopeExtension:
    return QwenImageScaleRopeExtension()