import gradio as gr
import re
from diffusers import StableDiffusionPipeline
import torch
from transformers import pipeline

# Carregue modelos (cache para performance)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    dtype=torch.float16 if device == "cuda" else torch.float32  # Corrigido: dtype em vez de torch_dtype
)
pipe = pipe.to(device)
# Removido: pipe.enable_model_cpu_offload()  # Isso causava o erro; opcional, mas não essencial aqui

extractor = pipeline("token-classification", model="valurank/MiniLM-L6-Keyword-Extraction")

def gerar_storyboard(roteiro, num_cenas=4, steps=20, guidance=7.5, estilos_customizados=""):
    # Passo 1: Divida o roteiro em cenas
    cenas = re.split(r'(?=CENA \d+:|Seção:|\n[A-Z][a-z]+:)', roteiro)[:num_cenas]  # Regex melhorada pra capturar seções como "A Invenção Invisível"
    cenas = [c.strip() for c in cenas if c.strip()]

    imagens = []
    for cena in cenas:
        # Passo 2: Extraia keywords
        keywords = [ent['word'] for ent in extractor(cena)[:5]]
        
        # Passo 3: Construa o prompt base
        prompt = f"Storyboard scene: {', '.join(keywords)}, cinematic, detailed, historical style"
        
        # Adicione estilos customizados (se fornecidos, senão usa vazio)
        if estilos_customizados.strip():
            prompt += f", {estilos_customizados}"
        
        print(f"Prompt gerado para cena: {prompt}")  # Log pra debug (remove se quiser)
        
        # Passo 4: Gere imagem
        with torch.autocast(device):
            img = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance).images[0]
        imagens.append(img)

    return imagens

# Interface Gradio com a nova Textbox
with gr.Blocks(title="Roteiro to Storyboard") as demo:
    gr.Markdown("# Gere Storyboards do Seu Roteiro com IA")
    
    with gr.Row():
        roteiro_input = gr.Textbox(label="Cole seu roteiro aqui", lines=10, placeholder="Ex: Máquina de Lavar – A Batalha Silenciosa...")
    
    with gr.Row():
        num_cenas = gr.Slider(2, 10, value=5, label="Número de Cenas")
        steps = gr.Slider(10, 50, value=15, label="Passos de Inferência (qualidade vs. velocidade)")  # Default menor pra acelerar
        guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale (criatividade)")
    
    # Nova Textbox para estilos
    estilos_input = gr.Textbox(
        label="Estilos Personalizados (opcional)", 
        lines=2, 
        placeholder="Ex: vintage washing machine, industrial revolution style, era vitoriana, vapor e engrenagens"
    )
    
    output = gr.Gallery(label="Storyboards Gerados")
    gerar_btn = gr.Button("Gerar Storyboard", variant="primary")
    
    # Conecte tudo à função
    gerar_btn.click(
        gerar_storyboard, 
        inputs=[roteiro_input, num_cenas, steps, guidance, estilos_input], 
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()