aduc_framework/managers/llama_scout_manager.py

# aduc_framework/managers/llama_scout_manager.py
#
# Versão 21.0.0 (Literal Function Replica)
# Implementa a função `generate_text_from_image` exatamente como no script
# da Meta para garantir 100% de fidelidade.

import yaml
import os
import logging
import torch
from PIL import Image
from typing import List, Callable
# --- IMPORTAÇÕES EXATAS DO SCRIPT DA META ---
from transformers import MllamaForConditionalGeneration, MllamaProcessor
from transformers import AutoTokenizer, AutoModelForCausalLM

from ..tools.hardware_manager import hardware_manager

logger = logging.getLogger(__name__)

# --- A FUNÇÃO ORIGINAL DA META, AGORA DENTRO DO NOSSO ARQUIVO ---
def generate_text_from_image(
    model, processor, image, prompt_text: str,
    temperature: float = 0.6, top_p: float = 0.9, max_new_tokens=2048
):
    """Generate text from image using model"""
    # 1. Garante que a imagem está no formato correto
    processed_image = image.convert("RGB")

    # 2. Constrói a estrutura da conversa
    conversation = [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]},
    ]

    # 3. Usa o `apply_chat_template` para criar o prompt
    prompt = processor.apply_chat_template(
        conversation, add_generation_prompt=True, tokenize=False
    )

    # 4. O processador combina a imagem e o texto
    inputs = processor(
        processed_image, prompt,
        text_kwargs={"add_special_tokens": False},
        return_tensors="pt"
    ).to(model.device)

    # 5. O modelo gera a resposta
    output = model.generate(
        **inputs,
        temperature=temperature,
        top_p=top_p,
        max_new_tokens=max_new_tokens,
        do_sample=True # Necessário para usar temperature e top_p
    )
    
    # 6. Decodifica e limpa a resposta
    full_response = processor.decode(output[0])
    clean_response = full_response[len(prompt):]

    if clean_response.endswith("<|eot_id|>"):
            clean_response = clean_response[:-len("<|eot_id|>")].strip()

    return clean_response.strip()
# -----------------------------------------------------------------

class LlamaScoutManager:
    """
    Uma casca fina que carrega os modelos e chama a função de inferência
    oficial da Meta.
    """
    def __init__(self, config: dict):
        self.hf_token = os.getenv("HF_TOKEN")
        if not self.hf_token: raise ValueError("HF_TOKEN é necessário.")

        multimodal_id = config['multimodal_model_id']
        helper_id = config['helper_model_id']

        logger.info(f"LLAMA SCOUT (Literal Replica): Carregando Cinegrafista: {multimodal_id}...")
        self.multimodal_processor = MllamaProcessor.from_pretrained(multimodal_id, token=self.hf_token)
        self.multimodal_model = MllamaForConditionalGeneration.from_pretrained(
            multimodal_id,
            torch_dtype=torch.bfloat16,
            use_safetensors=True,
            device_map="auto",
            token=self.hf_token
        )
        logger.info("LLAMA SCOUT (Literal Replica): Cinegrafista (Llama 3.2 Vision) carregado.")

        logger.info(f"LLAMA SCOUT (Literal Replica): Carregando Diretor: {helper_id}...")
        self.helper_tokenizer = AutoTokenizer.from_pretrained(helper_id, token=self.hf_token)
        self.helper_model = AutoModelForCausalLM.from_pretrained(
            helper_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            token=self.hf_token,
            attn_implementation="flash_attention_2"
        )
        logger.info("LLAMA SCOUT (Literal Replica): Diretor (Llama 3.1 8B) carregado.")

    @torch.inference_mode()
    def analyze_sequence(self, image_list: List[Image.Image], question: str, progress_callback: Callable = None) -> str:
        if not image_list: return "Nenhuma imagem fornecida."

        # A lógica agora é simples: chamar a função replicada.
        # Nós ainda usamos nossa lógica de "chunking" se necessário.
        if len(image_list) == 1:
            if progress_callback: progress_callback(0.2, f"Analisando 1 imagem com a função oficial...")
            return generate_text_from_image(self.multimodal_model, self.multimodal_processor, image_list[0], question)
        
        # Para múltiplas imagens, ainda precisamos da nossa lógica de orquestração.
        else:
            if progress_callback: progress_callback(0.1, f"Múltiplas imagens detectadas. Analisando em chunks...")
            # A API do modelo funciona melhor com uma imagem, então vamos analisar uma por uma
            # e depois pedir ao nosso Diretor para resumir.
            partial_analyses = []
            for i, image in enumerate(image_list):
                progress = 0.1 + (i / len(image_list)) * 0.8
                if progress_callback: progress_callback(progress, f"Analisando imagem {i+1}/{len(image_list)}...")
                
                # Criamos uma pergunta específica para cada imagem
                chunk_question = f"Esta é a imagem {i+1} de uma sequência. {question}"
                analysis = generate_text_from_image(self.multimodal_model, self.multimodal_processor, image, chunk_question)
                partial_analyses.append(analysis)
            
            return self._summarize_with_helper(partial_analyses, question, progress_callback)

    @torch.inference_mode()
    def _summarize_with_helper(self, partial_texts: List[str], original_question: str, progress_callback: Callable) -> str:
        if progress_callback: progress_callback(0.9, "Síntese com o Diretor 8B (Local)...")
        combined_partials = "\n\n---\n\n".join(f"Análise da Imagem {i+1}:\n{text}" for i, text in enumerate(partial_texts))
        prompt = (f"Você é um diretor de cinema. Sua visão é: '{original_question}'. "
                  f"Seu cinegrafista enviou os seguintes relatórios, um para cada imagem de uma sequência: {combined_partials}. "
                  "Sintetize esses relatórios em uma única resposta final, coesa e poderosa, "
                  "que atenda à sua visão original. Responda diretamente, sem mencionar os relatórios.")
        messages = [{"role": "user", "content": prompt}]
        input_ids = self.helper_tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(self.helper_model.device)
        outputs = self.helper_model.generate(input_ids, max_new_tokens=2048, do_sample=False)
        response = self.helper_tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
        return response.strip()

# (Placeholder e instanciação singleton permanecem iguais)
class LlamaScoutPlaceholder:
    def __init__(self, reason: str = "Motivo desconhecido"):
        logger.error(f"LlamaScoutManager não inicializado. Razão: {reason}. Placeholder em uso.")
        self.reason = reason
    def analyze_sequence(self, *args, **kwargs):
        return f"ERRO: Especialista Llama Scout indisponível. Razão: {self.reason}"

try:
    with open("config.yaml", 'r') as f: config = yaml.safe_load(f)
    llama_scout_config = config['specialists'].get('llama_scout')
    if llama_scout_config and llama_scout_config.get('gpus_required', 0) > 0:
        hardware_manager.allocate_gpus('LlamaScout', llama_scout_config['gpus_required'])
        llama_scout_manager_singleton = LlamaScoutManager(config=llama_scout_config)
        logger.info("Especialista de Análise (Original Recipe) pronto.")
    else:
        llama_scout_manager_singleton = LlamaScoutPlaceholder("Não habilitado na config.yaml")
except Exception as e:
    logger.critical(f"Falha CRÍTICA ao inicializar o LlamaScoutManager (Local): {e}", exc_info=True)
    llama_scout_manager_singleton = LlamaScoutPlaceholder(reason=str(e))