knishika62-ComfyUI-TextGenerateGemma3Prompt/nodes.py

import torch
import random

LTX23_SYSTEM_PROMPT_T2V = (
    "You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets. "
    "Your job is to create detailed prompts for the video bot. "
    "The prompts should be very descriptive and include details like the subject, setting, lighting, camera angle, and any other relevant details. "
    "You should also include details about the motion in the video. "
    "The prompts should be in English. "
    "Do not include any text in square brackets. "
    "Do not include any text that is not part of the prompt. "
    "Do not include any text that is not in English. "
    "Do not include any text that is not a prompt. "
    "The prompt should be a single paragraph. "
    "The prompt should be less than 200 words. "
    "The prompt should be a description of a video. "
    "The prompt should be in the present tense. "
    "The prompt should be in the third person. "
    "The prompt should not include any text that is not part of the video description."
)

LTX23_SYSTEM_PROMPT_I2V = (
    "You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets. "
    "Your job is to create detailed prompts for the video bot based on a reference image. "
    "The prompts should describe the continuation or animation of the scene shown in the image. "
    "Include details like the subject, setting, lighting, camera movement, and motion that would naturally follow from the image. "
    "The prompts should be in English. "
    "Do not include any text in square brackets. "
    "The prompt should be a single paragraph. "
    "The prompt should be less than 200 words. "
    "The prompt should be a description of a video that starts from the provided image. "
    "The prompt should be in the present tense. "
    "The prompt should be in the third person."
)

IMAGE_SOFT_TOKEN = "<image_soft_token>"


def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
    return (
        f"<start_of_turn>system\n{system_prompt}<end_of_turn>\n"
        f"<start_of_turn>user\n{user_prompt}<end_of_turn>\n"
        f"<start_of_turn>model\n"
    )


class TextGenerateGemma3Prompt:
    CATEGORY = "LTX-2.3"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)
    FUNCTION = "generate"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip": ("CLIP",),
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "mode": (["LTX-2.3", "Custom"],),
                "custom_system_prompt": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "You are a helpful assistant that generates detailed video prompts.",
                    },
                ),
                "max_length": ("INT", {"default": 256, "min": 1, "max": 4096}),
                "do_sample": ("BOOLEAN", {"default": True}),
                "temperature": (
                    "FLOAT",
                    {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01},
                ),
                "top_k": ("INT", {"default": 64, "min": 0, "max": 1000}),
                "top_p": (
                    "FLOAT",
                    {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
                "min_p": (
                    "FLOAT",
                    {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
                "repetition_penalty": (
                    "FLOAT",
                    {"default": 1.05, "min": 1.0, "max": 2.0, "step": 0.01},
                ),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
            },
            "optional": {
                "image": ("IMAGE",),
            },
        }

    def generate(
        self,
        clip,
        prompt,
        mode,
        custom_system_prompt,
        max_length,
        do_sample,
        temperature,
        top_k,
        top_p,
        min_p,
        repetition_penalty,
        seed,
        image=None,
    ):
        # Seed
        torch.manual_seed(seed)
        random.seed(seed)

        # Determine T2V or I2V
        is_i2v = image is not None

        # Select system prompt
        if mode == "LTX-2.3":
            system_prompt = LTX23_SYSTEM_PROMPT_I2V if is_i2v else LTX23_SYSTEM_PROMPT_T2V
        else:
            system_prompt = custom_system_prompt

        # Build user message
        user_message = prompt
        if is_i2v:
            user_message = f"{IMAGE_SOFT_TOKEN}\n{prompt}" if prompt else IMAGE_SOFT_TOKEN

        # Build full chat prompt
        full_prompt = build_chat_prompt(system_prompt, user_message)

        # Tokenize (pass image here so it gets embedded at <image_soft_token> positions)
        if is_i2v:
            tokens = clip.tokenize(full_prompt, image=image)
        else:
            tokens = clip.tokenize(full_prompt)

        # Generate
        gen_kwargs = {
            "max_length": max_length,
            "do_sample": do_sample,
            "temperature": temperature,
            "top_k": top_k,
            "top_p": top_p,
            "min_p": min_p,
            "repetition_penalty": repetition_penalty,
            "seed": seed,
        }
        output = clip.generate(tokens, **gen_kwargs)

        # Decode
        if hasattr(clip, "decode"):
            generated_text = clip.decode(output)
        else:
            generated_text = output

        # Strip the prompt prefix to return only generated portion
        if isinstance(generated_text, str) and generated_text.startswith(full_prompt):
            generated_text = generated_text[len(full_prompt):]

        # Clean up turn markers if present
        for marker in ["<end_of_turn>", "<start_of_turn>", "model\n"]:
            generated_text = generated_text.replace(marker, "")

        return (generated_text.strip(),)


NODE_CLASS_MAPPINGS = {
    "TextGenerateGemma3Prompt": TextGenerateGemma3Prompt,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TextGenerateGemma3Prompt": "Text Generate Gemma3 Prompt",
}