test.py

import re
from transformers import pipeline, PreTrainedTokenizerFast

from comfyui.nodes.node import Node
from typing import Union, Dict
from torch import Tensor

# Negative prompting utility function
def _build_negative_prompt(generated: str, input_prompt: str):
    # Remove any leaked model identifiers and action count from the output
    cleaned = re.sub(r'\b(model|version|gpt3\.5|gpt-3\.[0-9]+\|token|count|action)\b', '', generated)
    return f"Output is {input_prompt}, but it's NOT {cleaned}"

class LTX2MPSPromptArchitect(Node):
    SYSTEM_PROMPT = (
        "You are a master screenplay writer. Your job is to write scenes with vivid, "
        "realistic descriptions and natural-sounding dialogue when appropriate. Use proper grammar and punctuation. "
        "Follow the instructions given by the user in the 'User Direction' section of each prompt. "
        "Write in third person limited point of view. Do not include any meta-references or notes about yourself "
        "or your process — only describe what would be seen and heard in a film or TV show based on the scene provided."
    )

    # Node properties to be set in ComfyUI
    model_name: str = "gpt-3.5-turbo"  # OpenAI model ID for the GPT-3 variant used by default
    min_tokens: int = 0              # Minimum number of tokens to generate (default: 0)
    max_tokens: int = 100            # Maximum number of tokens to generate (default: 100)
    keep_model_loaded: bool = True   # Whether to keep the model in memory between calls

    def __init__(self):
        super().__init__()
        self.tokenizer = None
        self.model = None

    def load_model(self):
        if not self.keep_model_loaded:
            raise ValueError("Cannot manually load model when keep_model_loaded is False.")

        # Load the tokenizer and model from OpenAI's transformers library
        import transformers

        self.tokenizer = PreTrainedTokenizerFast.from_pretrained(self.model_name)
        self.model = pipeline("text-generation", model=self.model_name, tokenizer=self.tokenizer)

    def unload_model(self):
        if self.model is not None:
            del self.model
        if self.tokenizer is not None:
            del self.tokenizer

    def on_input(self, inputs: Dict[str, Union[str, Tensor]]):
        # Check if model is already loaded; if not, load it here
        if self.model is None:
            self.load_model()
        
        # Get the input text from ComfyUI
        user_input = inputs["user_input"]
        
        # Set default parameters for MPS optimization (adjust as needed)
        mps_batch_size = 16
        max_tokens_actual = self.max_tokens * 2  # Double the requested length for MPS optimization
        
        # Optional: allow a scene_context to be provided from a vision node (e.g., DALL-E 2)
        scene_context = inputs.get("scene_context", "")
        
        # Optional: allow custom LorA triggers to influence model behavior
        lora_triggers = inputs.get("lora_triggers", "")
        
        # Inject the scene context and LORA triggers into the input
        effective_input = (
            f"[SCENE CONTEXT FROM IMAGE — use this as the authoritative description "
            f"of the subject and setting; do not invent or contradict it]\n"
            f"{scene_context.strip()}" + "\n" + (lora_triggers.strip() if lora_triggers else "") + "\n\n"
            f"[USER DIRECTION — apply this as action, style, and mood over the above scene]\n"
            f"{user_input}"
        )
        
        # Generate text using the MPS-optimized model
        with torch.no_grad():
            output = self.model(prompt=effective_input, max_length=max_tokens_actual, min_length=self.min_tokens, num_return_sequences=1)
        
        # Decode the generated tokens into plain text and remove any trailing newline or whitespace
        result = self.tokenizer.decode(output[0], skip_special_tokens=True).strip()

        # Regex clean as a last-resort safety net (should rarely trigger now)
        neg_prompt = self._clean_output(result, user_input)
        
        return {
            "generated": result,
            "negative_prompt": neg_prompt,
        }


    @staticmethod
    def _clean_output(generated: str, input_prompt: str) -> str:
        # Remove any leaked model identifiers and action count from the output
        cleaned = re.sub(r'\b(model|version|gpt3\.5-turbo|token|count|action)\b', '', generated, flags=re.IGNORECASE).strip()
        
        # Check if negative prompting is needed (anything that looks like a model response)
        bullshit_markers = ["[", "]", "(", ")", "•", ";"]
        if any(x in cleaned for x in bullshit_markers) or len(cleaned.split()) > 5:
            return _build_negative_prompt(generated, input_prompt)
        
        # No negative prompt needed — just return the clean text
        return generated


        
    
# Node definitions for ComfyUI
LTX2MPS_NODE = {
    "name": "LTX2MPS Easy Prompt [MPS Optimized] By blackest",
    "description": "An optimized version of OpenAI's GPT-3.5-turbo for generating text with MPS support.",
    "category": "Text",
    "icon": "text_fields.png",
    "inputs": {
        "user_input": {
            "name": "User Direction",
            "description": "The action, style, and mood for the AI to follow in the scene.",
            "type": "string",
            "required": True
        },
        "scene_context": {
            "name": "Scene Context (Optional)",
            "description": "Provide a description of the scene from an image if available, "
                           "to be used as the initial context.",
            "type": "string",
            "required": False
        },
        "lora_triggers": {
            "name": "LORA Triggers (Optional)",
            "description": "Custom LorA triggers to influence model behavior.",
            "type": "string",
            "required": False
        }
    },
    "outputs": {
        "generated": {
            "name": "Generated Text",
            "description": "The text generated by the GPT-3.5-turbo model.",
            "type": "string"
        },
        "negative_prompt": {
            "name": "Negative Prompt",
            "description": "A negative version of the output to be used for prompting the model again.",
            "type": "string"
        }
    },
    "classes": [
        LTX2MPSPromptArchitect,
        LTX2MPSUnloadModel
    ],
}