import gradio as gr
import gc
import torch
import spaces  # <-- A biblioteca do ZeroGPU está de volta!
import os
from deep_translator import GoogleTranslator
from transformers import (
    BlipProcessor, BlipForConditionalGeneration,
    VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer,
    AutoProcessor, AutoModelForCausalLM, 
    LlavaForConditionalGeneration, PaliGemmaForConditionalGeneration,
    Qwen2VLForConditionalGeneration
)

def traduzir_para_pt(texto):
    try:
        return GoogleTranslator(source='en', target='pt').translate(texto)
    except Exception as e:
        return f"{texto} (Erro na tradução: {str(e)})"

# O decorador do ZeroGPU ativado para pedir a placa de vídeo ao servidor
@spaces.GPU(duration=120)
def analisar_imagem(imagem, tamanho):
    if imagem is None:
        return ["Nenhuma imagem carregada."] * 6
    
    resultados = []
    hf_token = os.environ.get("HF_TOKEN")

    # Radar de hardware: detecta no milissegundo do clique se a GPU foi liberada
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype_modernos = torch.float16 if device == "cuda" else torch.bfloat16

    if "Pequeno" in tamanho:
        max_tok = 40
        prompt_vlms = "Describe this image briefly in one or two sentences."
    elif "Médio" in tamanho:
        max_tok = 120
        prompt_vlms = "Describe this image in detail."
    else: 
        max_tok = 300
        prompt_vlms = "Describe every element of this image in extensive and exhaustive detail."
        
    def limpar_memoria():
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()
            
    limpar_memoria()
    
    # =========================================================
    # GERAÇÃO CLÁSSICA (2021 - 2022)
    # =========================================================
    
    try:
        print(f"\n[ HARDWARE ZEROGPU ATIVO: {device.upper()} ]")
        print("Carregando BLIP...")
        processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
        model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
        inputs = processor(imagem, return_tensors="pt").to(device)
        out = model.generate(**inputs, max_new_tokens=max_tok)
        resultados.append(traduzir_para_pt(processor.decode(out[0], skip_special_tokens=True)))
        del processor, model, inputs, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro BLIP: {str(e)}")
        
    try:
        print("Carregando ViT-GPT2...")
        model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning").to(device)
        feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
        tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
        pixel_values = feature_extractor(images=imagem, return_tensors="pt").pixel_values.to(device)
        out = model.generate(pixel_values, max_new_tokens=max_tok, num_beams=4)
        resultados.append(traduzir_para_pt(tokenizer.decode(out[0], skip_special_tokens=True)))
        del model, feature_extractor, tokenizer, pixel_values, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro ViT-GPT2: {str(e)}")

    try:
        print("Carregando GIT...")
        processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
        model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco").to(device)
        pixel_values = processor(images=imagem, return_tensors="pt").pixel_values.to(device)
        out = model.generate(pixel_values=pixel_values, max_new_tokens=max_tok)
        resultados.append(traduzir_para_pt(processor.decode(out[0], skip_special_tokens=True)))
        del processor, model, pixel_values, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro GIT: {str(e)}")

    # =========================================================
    # GERAÇÃO MODERNA (VLMs de 2023 - 2024)
    # =========================================================

    try:
        print("Carregando Qwen2-VL...")
        model_id = "Qwen/Qwen2-VL-2B-Instruct"
        model = Qwen2VLForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype_modernos).to(device)
        processor = AutoProcessor.from_pretrained(model_id)
        
        messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_vlms}]}]
        text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = processor(text=[text_prompt], images=[imagem], padding=True, return_tensors="pt").to(device)
        
        out = model.generate(**inputs, max_new_tokens=max_tok)
        generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, out)]
        texto_en = processor.decode(generated_ids_trimmed[0], skip_special_tokens=True)
        
        resultados.append(traduzir_para_pt(texto_en))
        del model, processor, inputs, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro Qwen2-VL: {str(e)}")
        
    try:
        print("Carregando PaliGemma...")
        model_id = "google/paligemma-3b-mix-224"
        processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype_modernos, token=hf_token).to(device)
        inputs = processor(text="caption", images=imagem, return_tensors="pt").to(device)
        out = model.generate(**inputs, max_new_tokens=max_tok)
        resultados.append(traduzir_para_pt(processor.decode(out[0], skip_special_tokens=True)))
        del processor, model, inputs, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro PaliGemma: {str(e)}")

    try:
        print("Carregando LLaVA...")
        model_id = "llava-hf/llava-1.5-7b-hf"
        processor = AutoProcessor.from_pretrained(model_id)
        model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype_modernos).to(device)
        prompt = f"USER: <image>\n{prompt_vlms}\nASSISTANT:"
        inputs = processor(images=imagem, text=prompt, return_tensors="pt").to(device)
        out = model.generate(**inputs, max_new_tokens=max_tok)
        texto_gerado = processor.decode(out[0], skip_special_tokens=True)
        texto_en = texto_gerado.split("ASSISTANT:")[-1].strip()
        resultados.append(traduzir_para_pt(texto_en))
        del processor, model, inputs, out
        limpar_memoria()
    except Exception as e: resultados.append(f"Erro LLaVA: {str(e)}")

    return tuple(resultados)

with gr.Blocks() as demo:
    gr.Markdown("# 👁️ Linha do Tempo: Modelos Visuais")
    gr.Markdown("Faça o upload de uma imagem, selecione a complexidade da análise e observe o salto tecnológico. **Sistema ZeroGPU Ativo.**")
    
    with gr.Row():
        with gr.Column(scale=1):
            entrada_imagem = gr.Image(type="pil", label="Imagem de Entrada")
            seletor_tamanho = gr.Radio(
                choices=["Pequeno (aprox. 3 linhas)", "Médio (aprox. 6 linhas)", "Grande (aprox. 10 linhas)"],
                value="Médio (aprox. 6 linhas)",
                label="Extensão da Descrição"
            )
            botao_analisar = gr.Button("Processar nos 6 Modelos", variant="primary")
        
        with gr.Column(scale=1):
            gr.Markdown("### Geração Clássica")
            saida_blip = gr.Textbox(label="BLIP (Salesforce - Fev/2022)", lines=4)
            saida_vit = gr.Textbox(label="ViT-GPT2 (NLP Connect - 2021)", lines=4)
            saida_git = gr.Textbox(label="GIT (Microsoft - Mai/2022)", lines=4)

        with gr.Column(scale=1):
            gr.Markdown("### Geração Moderna (VLMs)")
            saida_qwen = gr.Textbox(label="Qwen2-VL 2B (Alibaba - Set/2024)", lines=4)
            saida_pali = gr.Textbox(label="PaliGemma 3B (Google - Mai/2024)", lines=4)
            saida_llava = gr.Textbox(label="LLaVA 1.5 7B (Out/2023)", lines=4)
            
    botao_analisar.click(
        fn=analisar_imagem,
        inputs=[entrada_imagem, seletor_tamanho],
        outputs=[saida_blip, saida_vit, saida_git, saida_qwen, saida_pali, saida_llava],
        api_name="processar",
        show_progress="full"
    )

if __name__ == "__main__":
    # Mantém a fila (queue) ativa para evitar o limite de tempo do navegador
    demo.queue(api_open=False)
    demo.launch(theme=gr.themes.Base())