import gradio as gr
from transformers import (
    BlipProcessor, BlipForConditionalGeneration,
    WhisperProcessor, WhisperForConditionalGeneration,
    AutoTokenizer, AutoModelForCausalLM
)
from diffusers import StableDiffusionPipeline
import torch
import librosa
import subprocess
from gtts import gTTS
import os

# Carregar modelos
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

sd_pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16).to("cuda")

whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")

# Modelo de chatbot (OpenAssistant em português)
tokenizer = AutoTokenizer.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5")
chatbot_model = AutoModelForCausalLM.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5")

# Funções principais
def gerar_imagem(prompt):
    imagem = sd_pipe(prompt).images[0]
    return imagem

def descrever_imagem(imagem):
    inputs = blip_processor(imagem, return_tensors="pt")
    out = blip_model.generate(**inputs)
    return blip_processor.decode(out[0], skip_special_tokens=True)

def transcrever_audio(audio):
    speech, _ = librosa.load(audio, sr=16000)
    input_features = whisper_processor(speech, sampling_rate=16000, return_tensors="pt").input_features
    predicted_ids = whisper_model.generate(input_features)
    return whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]

def responder_chat(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = chatbot_model.generate(**inputs, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def executar_comando(comando):
    try:
        resultado = subprocess.run(comando, shell=True, capture_output=True, text=True)
        return resultado.stdout or "Comando executado!"
    except Exception as e:
        return f"Erro: {str(e)}"

def texto_para_voz(texto):
    tts = gTTS(texto, lang='pt-br')
    tts.save("resposta.mp3")
    return "resposta.mp3"

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 IA Multimodal Completa")
    
    with gr.Tab("Gerar Imagem"):
        texto_imagem = gr.Textbox(label="Descreva a imagem que deseja gerar")
        saida_imagem = gr.Image(label="Imagem Gerada")
        btn_imagem = gr.Button("Gerar")
    
    with gr.Tab("Descrever Imagem"):
        entrada_imagem = gr.Image(label="Envie uma imagem", type="pil")
        texto_descricao = gr.Textbox(label="Descrição Gerada")
        btn_descricao = gr.Button("Descrever")
    
    with gr.Tab("Transcrever Áudio"):
        entrada_audio = gr.Audio(label="Grave ou envie um áudio", type="filepath")
        texto_transcricao = gr.Textbox(label="Transcrição")
        btn_transcricao = gr.Button("Transcrever")
    
    with gr.Tab("Chatbot"):
        chatbot_input = gr.Textbox(label="Pergunte algo")
        chatbot_output = gr.Textbox(label="Resposta")
        chatbot_audio = gr.Audio(label="Resposta em Áudio", visible=True)
        chatbot_btn = gr.Button("Enviar")
    
    with gr.Tab("Executar Comandos"):
        comando_input = gr.Textbox(label="Comando (ex: pip install numpy)")
        comando_output = gr.Textbox(label="Resultado")
        comando_btn = gr.Button("Executar")
    
    # Conexões
    btn_imagem.click(gerar_imagem, inputs=texto_imagem, outputs=saida_imagem)
    btn_descricao.click(descrever_imagem, inputs=entrada_imagem, outputs=texto_descricao)
    btn_transcricao.click(transcrever_audio, inputs=entrada_audio, outputs=texto_transcricao)
    chatbot_btn.click(responder_chat, inputs=chatbot_input, outputs=chatbot_output)
    chatbot_btn.click(texto_para_voz, inputs=chatbot_output, outputs=chatbot_audio)
    comando_btn.click(executar_comando, inputs=comando_input, outputs=comando_output)

demo.launch()