managers/prompt_enhancer_manager.py

# managers/prompt_enhancer_manager.py
#
# Copyright (C) 2025 Carlos Rodrigues dos Santos
#
# Version: 1.0.1
#
# This version fixes an AttributeError during the loading of the Florence-2 model
# by explicitly setting the attention implementation to "eager". This bypasses an
# incompatibility between a newer version of the transformers library and the
# model's specific code, ensuring robust initialization.

import torch
import logging
import yaml
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM, AutoTokenizer
from pathlib import Path

logger = logging.getLogger(__name__)

# O prompt de sistema que guiará nosso LLM
ENHANCER_SYSTEM_PROMPT = """You are an expert cinematic director. Your task is to write a single, rich, cinematic motion prompt.
Analyze the user's goal and the provided image caption. Synthesize them into a flowing, descriptive paragraph under 150 words.
Focus on the action, character expressions, camera movement, and environment. Start directly with the action.
The final prompt must be a direct instruction for a video generation AI."""

class PromptEnhancerManager:
    def __init__(self):
        logger.info("Initializing Prompt Enhancer Manager...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32

        try:
            with open("config.yaml", 'r') as f:
                config = yaml.safe_load(f)['specialists']['prompt_enhancer']
            
            caption_model_name = config['image_caption_model']
            llm_model_name = config['llm_model']
            
            logger.info(f"Loading Image Caption Model: {caption_model_name}...")
            self.caption_processor = AutoProcessor.from_pretrained(caption_model_name, trust_remote_code=True)
            
            # <--- CORREÇÃO AQUI --->
            # Forçamos o uso da implementação de atenção "eager" (a mais antiga e compatível)
            # para evitar o erro de atributo _supports_sdpa.
            self.caption_model = AutoModelForCausalLM.from_pretrained(
                caption_model_name, 
                torch_dtype=self.dtype, 
                trust_remote_code=True,
                attn_implementation="eager"  # Adiciona este parâmetro
            ).to(self.device)
            # <--- FIM DA CORREÇÃO --->
            
            logger.info(f"Loading LLM for Prompt Enhancement: {llm_model_name}...")
            self.llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
            self.llm_model = AutoModelForCausalLM.from_pretrained(
                llm_model_name,
                torch_dtype=self.dtype,
                device_map="auto" 
            )
            
            logger.info("Prompt Enhancer Manager initialized successfully.")
        except Exception as e:
            logger.critical("Failed to initialize PromptEnhancerManager.", exc_info=True)
            raise e

    @torch.no_grad()
    def generate_enhanced_prompt(self, image: Image.Image, user_prompt: str) -> str:
        # ... (o resto da classe permanece exatamente o mesmo) ...
        logger.info("Generating enhanced prompt...")
        
        caption_task_prompt = "<MORE_DETAILED_CAPTION>"
        inputs = self.caption_processor(
            text=caption_task_prompt, images=image, return_tensors="pt"
        ).to(self.device, self.dtype)
        
        generated_ids = self.caption_model.generate(**inputs, max_new_tokens=1024)
        generated_texts = self.caption_processor.batch_decode(generated_ids, skip_special_tokens=True)
        # O parse da resposta do Florence-2 pode variar, vamos torná-lo mais robusto
        parsed_caption = generated_texts[0].replace(caption_task_prompt, "").strip()
        image_caption = parsed_caption
        logger.info(f"Generated Image Caption: '{image_caption}'")

        messages = [
            {"role": "system", "content": ENHANCER_SYSTEM_PROMPT},
            {"role": "user", "content": f"My Goal: '{user_prompt}'\n\nReference Image Scene: '{image_caption}'"}
        ]
        
        input_ids = self.llm_tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        ).to(self.llm_model.device)
        
        outputs = self.llm_model.generate(input_ids, max_new_tokens=256, do_sample=True, temperature=0.6, top_p=0.9)
        response = self.llm_tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
        
        logger.info(f"LLM Enhanced Prompt: '{response}'")
        return response.strip()

# --- Singleton Instantiation ---
try:
    prompt_enhancer_manager_singleton = PromptEnhancerManager()
except Exception as e:
    prompt_enhancer_manager_singleton = None
    raise e