Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import re
 from diffusers import StableDiffusionPipeline
 import torch
-from transformers import pipeline  # Para extração de keywords
+from transformers import pipeline
 
 # Carregue modelos (cache para performance)
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -11,39 +11,64 @@ pipe = StableDiffusionPipeline.from_pretrained(
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
 )
 pipe = pipe.to(device)
+pipe.enable_model_cpu_offload()  # Otimização pra acelerar e economizar memória
 
 extractor = pipeline("token-classification", model="valurank/MiniLM-L6-Keyword-Extraction")
 
-def gerar_storyboard(roteiro, num_cenas=4, steps=20, guidance=7.5):
+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:)', roteiro)[:num_cenas]  # Ajuste para seu formato
+    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()]
 
-    # Passo 2: Extraia prompts visuais por cena
     imagens = []
     for cena in cenas:
+        # Passo 2: Extraia keywords
         keywords = [ent['word'] for ent in extractor(cena)[:5]]
-        prompt = f"Storyboard scene: {', '.join(keywords)}, cinematic, detailed, historical style"  # Adapte para tema (ex.: "máquina de lavar antiga")
         
-        # Passo 3: Gere imagem
+        # 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  # Retorna lista de PIL images para galeria
+    return imagens
 
-# Interface Gradio
+# Interface Gradio com a nova Textbox
 with gr.Blocks(title="Roteiro to Storyboard") as demo:
     gr.Markdown("# Gere Storyboards do Seu Roteiro com IA")
-    roteiro_input = gr.Textbox(label="Cole seu roteiro aqui", lines=10, placeholder="Ex: Máquina de Lavar – A Batalha Silenciosa...")
-    num_cenas = gr.Slider(2, 10, value=5, label="Número de Cenas")
-    steps = gr.Slider(10, 50, value=20, label="Passos de Inferência (qualidade vs. velocidade)")
-    guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale (criatividade)")
+    
+    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")
+    gerar_btn = gr.Button("Gerar Storyboard", variant="primary")
     
-    gerar_btn.click(gerar_storyboard, inputs=[roteiro_input, num_cenas, steps, guidance], outputs=output)
+    # 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()