LTXVPromptEnhancer Issue Fix

prompt_enhancer_nodes-fix.py

@comfy_node(name="LTXVPromptEnhancer")
class LTXVPromptEnhancer:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "prompt": ("STRING",),
                "prompt_enhancer": ("LTXV_PROMPT_ENHANCER",),
                "max_resulting_tokens": (
                    "INT",
                    {"default": 256, "min": 32, "max": 512},
                ),
            },
            "optional": {
                "image_prompt": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("str",)
    FUNCTION = "enhance"
    CATEGORY = "lightricks/LTXV"
    TITLE = "LTXV Prompt Enhancer"
    OUTPUT_NODE = False
    DESCRIPTION = (
        "Enhances text prompts for image generation using LLMs. "
        "Optionally incorporates reference images to create more contextually relevant descriptions."
    )

    def enhance(
        self,
        prompt,
        prompt_enhancer: comfy.model_patcher.ModelPatcher,
        image_prompt: torch.Tensor = None,
        max_resulting_tokens=256,
    ):
        comfy.model_management.free_memory(
            prompt_enhancer.memory_required([]),
            comfy.model_management.get_torch_device(),
        )

        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        model = prompt_enhancer.model.to(device)

        image_conditioning = None
        if image_prompt is not None:
            permuted_image = image_prompt.permute(3, 0, 1, 2)[None, :].to(device)
            image_conditioning = [(permuted_image, 0, 1.0)]

        enhanced_prompt = model(prompt, image_conditioning, max_resulting_tokens)
        return (enhanced_prompt[0],)