app.py

import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random
import os
import json
import io
import uuid
from gradio_client import Client as client_gradio
from supabase import create_client, Client
from datetime import datetime
import requests
import logging

# Configuração de logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Inicializa Supabase
url: str = os.getenv('SUPABASE_URL')
key: str = os.getenv('SUPABASE_KEY')
supabase: Client = create_client(url, key)

# Obtém token da Hugging Face
hf_token = os.getenv("HF_TOKEN")

# Inicializa o modelo base FLUX.1-dev
base_model = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(
    base_model, 
    torch_dtype=torch.float16,
    use_safetensors=True
)

# Move o modelo para GPU
pipe.to("cuda")

# Definição dos LoRA e Trigger Words - Mais eficazes para gerar os personagens corretamente
lora_models = {
    "Paula": {
        "repo": "vcollos/Paula2",
        "weights": "Paula P.safetensors",
        "trigger_word": "woman with long blonde hair named Paula",
        "negative_word": "men, male, man, masculine features, dark hair",
        "character_desc": "a beautiful woman with long blonde hair, feminine features, soft facial features"
    },
    "Vivi": {
        "repo": "vcollos/Vivi",
        "weights": "Vivi.safetensors",
        "trigger_word": "man with dark hair named Vivi",
        "negative_word": "women, female, woman, feminine features, blonde hair",
        "character_desc": "a handsome man with dark hair, masculine features, defined jawline"
    }
}

# Carrega os LoRAs disponíveis
for name, details in lora_models.items():
    try:
        pipe.load_lora_weights(details["repo"], weight_name=details["weights"], adapter_name=name)
        logger.info(f"✅ LoRA {name} carregado")
    except Exception as e:
        logger.error(f"❌ Erro ao carregar o LoRA {name}: {e}")

# Define seed máximo
MAX_SEED = 2**32 - 1

def upload_image_to_supabase(image, filename):
    """ Faz upload da imagem para o Supabase Storage e retorna a URL pública. """
    img_bytes = io.BytesIO()
    image.save(img_bytes, format="PNG")
    img_bytes.seek(0)  # Move para o início do arquivo

    storage_path = f"images/{filename}"

    try:
        # Faz upload da imagem para o Supabase
        supabase.storage.from_("images").upload(storage_path, img_bytes.getvalue(), {"content-type": "image/png"})

        # Retorna a URL pública da imagem
        base_url = f"{url}/storage/v1/object/public/images"
        return f"{base_url}/{storage_path}"
    except Exception as e:
        logger.error(f"❌ Erro no upload da imagem: {e}")
        return None

# Dicionário simples de tradução português-inglês para palavras comuns
translations = {
    "homem": "man",
    "mulher": "woman",
    "juntos": "together",
    "e": "and",
    "com": "with",
    "dois": "two",
    "duas": "two",
    "pessoas": "people",
    "pessoa": "person",
    "sentado": "sitting",
    "sentada": "sitting",
    "em pé": "standing",
    "conversa": "conversation",
    "conversando": "talking",
    "falando": "talking",
    "praia": "beach",
    "jardim": "garden",
    "casa": "house",
    "cidade": "city",
    "parque": "park",
    "floresta": "forest",
    "montanha": "mountain",
    "rio": "river",
    "lago": "lake",
    "mar": "sea",
    "oceano": "ocean",
    "céu": "sky",
    "nuvem": "cloud",
    "sol": "sun",
    "lua": "moon",
    "estrela": "star",
    "dia": "day",
    "noite": "night",
    "manhã": "morning",
    "tarde": "afternoon",
    "amigo": "friend",
    "amiga": "friend",
    "casal": "couple",
    "família": "family",
    "irmão": "brother",
    "irmã": "sister",
    "pai": "father",
    "mãe": "mother",
    "filho": "son",
    "filha": "daughter",
    "avô": "grandfather",
    "avó": "grandmother",
    "tio": "uncle",
    "tia": "aunt",
    "primo": "cousin",
    "prima": "cousin",
    "namorado": "boyfriend",
    "namorada": "girlfriend",
    "marido": "husband",
    "esposa": "wife",
    "amor": "love",
    "feliz": "happy",
    "triste": "sad",
    "bravo": "angry",
    "assustado": "scared",
    "surpreso": "surprised",
    "cansado": "tired",
    "entediado": "bored",
    "excitado": "excited",
    "confuso": "confused",
    # Adiciona mais termos de ambientes/ações/situações
    "café": "coffee shop",
    "restaurante": "restaurant",
    "cinema": "movie theater",
    "shopping": "mall",
    "biblioteca": "library",
    "escritório": "office",
    "hotel": "hotel",
    "aeroporto": "airport",
    "estação": "station",
    "hospital": "hospital",
    "escola": "school",
    "universidade": "university",
    "igreja": "church",
    "teatro": "theater",
    "museu": "museum",
    "bar": "bar",
    "festa": "party",
    "casamento": "wedding",
    "aniversário": "birthday",
    "caminhando": "walking",
    "correndo": "running",
    "dançando": "dancing",
    "cantando": "singing",
    "tocando": "playing",
    "dirigindo": "driving",
    "nadando": "swimming",
    "assistindo": "watching",
    "lendo": "reading",
    "escrevendo": "writing",
    "cozinhando": "cooking",
    "comendo": "eating",
    "bebendo": "drinking",
    "dormindo": "sleeping",
    "trabalhando": "working",
    "estudando": "studying",
    "fotografando": "photographing",
    "pintando": "painting",
    "desenhando": "drawing"
}

def simple_translate(text):
    """
    Função simples para traduzir texto para inglês usando o dicionário de traduções.
    """
    translated_text = text.lower()
    
    for pt, en in translations.items():
        # Use word boundary to avoid partial matches
        translated_text = translated_text.replace(f" {pt} ", f" {en} ")
        translated_text = translated_text.replace(f" {pt},", f" {en},")
        translated_text = translated_text.replace(f" {pt}.", f" {en}.")
        
        # Check for word at beginning of string
        if translated_text.startswith(f"{pt} "):
            translated_text = f"{en} " + translated_text[len(pt)+1:]
        
        # Check for word at end of string
        if translated_text.endswith(f" {pt}"):
            translated_text = translated_text[:-len(pt)-1] + f" {en}"
    
    logger.info(f"Texto traduzido: {translated_text}")
    return translated_text

@spaces.GPU(duration=80)
def run_lora(
    prompt, cfg_scale, steps, randomize_seed, seed, width, height, 
    lora_option, lora_scale_1, lora_scale_2, lora_balance, 
    translate_prompt, use_negative_prompt, quality_preset,
    progress=gr.Progress(track_tqdm=True)
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    original_prompt = prompt  # Guarda o prompt original para metadados
    
    # Traduz o prompt se a opção estiver ativada
    if translate_prompt:
        prompt = simple_translate(prompt)
    
    # Trunca o prompt para 77 tokens para evitar erro do CLIP
    prompt_tokens = prompt.split()[:77]
    prompt = " ".join(prompt_tokens)

    # Define qual LoRA usar com base na seleção do usuário
    selected_loras = []
    adapter_weights = []
    negative_prompt = ""
    
    # Aplica preset de qualidade
    if quality_preset == "Alta Qualidade":
        quality_terms = ", professional photography, detailed, high quality, 8k, masterpiece, best quality"
    elif quality_preset == "Artístico":
        quality_terms = ", artistic, cinematic lighting, dramatic, professional, detailed"
    elif quality_preset == "Realista":
        quality_terms = ", photorealistic, detailed skin, detailed face, high detail, realistic"
    else:  # Nenhum
        quality_terms = ""
    
    # Modificado para melhorar a caracterização de personagens
    if lora_option == "Paula":
        selected_loras.append("Paula")
        adapter_weights.append(lora_scale_1)
        # Adiciona mais ênfase no personagem e suas características
        prompt = f"{lora_models['Paula']['trigger_word']}, {lora_models['Paula']['character_desc']}, {prompt}{quality_terms}"
        if use_negative_prompt:
            negative_prompt = lora_models['Paula']['negative_word']
        
    elif lora_option == "Vivi":
        selected_loras.append("Vivi")
        adapter_weights.append(lora_scale_2)
        # Adiciona mais ênfase no personagem e suas características
        prompt = f"{lora_models['Vivi']['trigger_word']}, {lora_models['Vivi']['character_desc']}, {prompt}{quality_terms}"
        if use_negative_prompt:
            negative_prompt = lora_models['Vivi']['negative_word']
        
    elif lora_option == "Ambos":
        # Usa o balance slider para ajustar a proporção entre os dois LoRAs
        p_weight = lora_scale_1 * lora_balance
        v_weight = lora_scale_2 * (2 - lora_balance)
        
        selected_loras = ["Paula", "Vivi"]
        adapter_weights = [p_weight, v_weight]
        
        # Quando usando ambos, adiciona descrições mais explícitas para cada personagem
        prompt = f"{lora_models['Paula']['trigger_word']} and {lora_models['Vivi']['trigger_word']} together, side by side, a blonde woman and a dark-haired man, {prompt}{quality_terms}"

    pipe.set_adapters(selected_loras, adapter_weights)
    
    # Adiciona log para depuração
    logger.info(f"Prompt Final: {prompt}")
    logger.info(f"Negative Prompt: {negative_prompt}")
    logger.info(f"LoRA selecionado: {lora_option}, Pesos: {adapter_weights}")

    # Gera a imagem com precisão de 16 bits
    with torch.autocast("cuda"):
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator
        ).images[0]

    # Define um nome único para a imagem
    filename = f"image_{seed}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.png"

    try:
        image_url = upload_image_to_supabase(image, filename)
        if image_url:
            logger.info(f"✅ Imagem salva no Supabase: {image_url}")
        else:
            logger.error("❌ Erro: URL da imagem retornou None")
            return image, seed, prompt
    except Exception as e:
        logger.error(f"❌ Erro ao fazer upload da imagem: {e}")
        return image, seed, prompt

    # Salva todos os metadados no Supabase
    try:
        response = supabase.table("images").insert({
            "prompt": original_prompt,  # Salva o prompt original
            "translated_prompt": prompt if translate_prompt else original_prompt,
            "full_prompt": prompt,      # Salva o prompt completo com trigger words
            "negative_prompt": negative_prompt,
            "quality_preset": quality_preset,
            "cfg_scale": cfg_scale,
            "steps": steps,
            "seed": seed,
            "lora_option": lora_option,
            "lora_scale_1": lora_scale_1,
            "lora_scale_2": lora_scale_2,
            "lora_balance": lora_balance,
            "image_url": image_url,
            "created_at": datetime.utcnow().isoformat()
        }).execute()
        
        if response.data:
            logger.info("✅ Metadados salvos no Supabase")
        else:
            logger.error("❌ Erro: Resposta vazia do Supabase")
        
    except Exception as e:
        logger.error(f"❌ Erro ao salvar metadados no Supabase: {e}")

    return image, seed, prompt

# Interface Gradio
gr_theme = os.getenv("THEME")
with gr.Blocks(theme=gr_theme) as app:
    gr.Markdown("# Paula & Vivi Image Generator")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.TextArea(
                label="Prompt", 
                placeholder="Digite um prompt descrevendo o ambiente e ações (pode ser em português)", 
                lines=3
            )
            generate_button = gr.Button("Gerar Imagem", variant="primary")
            
            with gr.Accordion("Configurações Básicas", open=True):
                translate_prompt = gr.Checkbox(True, label="Traduzir prompt do português para inglês")
                quality_preset = gr.Radio(
                    ["Nenhum", "Alta Qualidade", "Artístico", "Realista"], 
                    label="Preset de Qualidade", 
                    value="Alta Qualidade"
                )
                cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=8.0)
                steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=35)
                
            with gr.Accordion("Tamanho e Seed", open=False):
                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=768)  
                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=1024)
                randomize_seed = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=556215326)
            
            with gr.Accordion("Configurações de LoRA", open=True):
                lora_option = gr.Radio(
                    ["Paula", "Vivi", "Ambos"],
                    label="Escolha o LoRA", 
                    value="Ambos"
                )
                use_negative_prompt = gr.Checkbox(
                    True, 
                    label="Usar negative prompt", 
                    info="Ajuda a evitar mistura de características"
                )
                with gr.Group(visible=True) as lora_controls:
                    lora_scale_1 = gr.Slider(
                        label="Intensidade do LoRA (Paula)", 
                        minimum=0, maximum=1, step=0.05, value=0.85
                    )
                    lora_scale_2 = gr.Slider(
                        label="Intensidade do LoRA (Vivi)", 
                        minimum=0, maximum=1, step=0.05, value=0.85
                    )
                    lora_balance = gr.Slider(
                        label="Balanço entre personagens", 
                        minimum=0.5, maximum=1.5, step=0.05, value=1.0, 
                        info="Valores acima de 1.0 favorecem Paula, abaixo de 1.0 favorecem Vivi"
                    )

        with gr.Column(scale=2):
            result = gr.Image(label="Imagem Gerada", type="pil")
            final_prompt = gr.Textbox(label="Prompt Final", lines=3)
            
            with gr.Accordion("Instruções", open=True):
                gr.Markdown("""
                ### Como escrever prompts eficientes:
                
                **O que colocar no prompt:**
                - **Ambiente/cenário**: "em uma praia", "em um café", "em uma floresta"
                - **Ações/atividades**: "conversando", "caminhando", "segurando mãos"
                - **Roupas/acessórios**: "vestido azul", "terno preto", "chapéu de palha"
                - **Iluminação/hora do dia**: "pôr do sol", "luz noturna", "iluminação suave"
                
                **O que NÃO precisa incluir:**
                - Não mencione "Paula" ou "Vivi" - o sistema já adiciona isso
                - Não mencione "mulher loira" ou "homem moreno" - isso já está incluído
                
                **Exemplos de bons prompts:**
                - "Em um café à beira-mar, conversando e rindo, pôr do sol"
                - "Caminhando em um parque de outono, roupas elegantes"
                - "Em uma festa, dançando juntos, luzes coloridas"
                """)
            
            with gr.Accordion("Dicas para melhores resultados", open=False):
                gr.Markdown("""
                ### Para um personagem só:
                - Mantenha o CFG Scale alto (7-9)
                - Intensidade do LoRA em 0.85-0.95
                - Deixe "Usar negative prompt" ativado
                - Use o preset "Realista" para fotos mais realistas
                
                ### Para ambos personagens juntos:
                - Use valores iguais para intensidade dos LoRAs
                - Mencione explicitamente que estão "juntos" ou "lado a lado"
                - Um CFG Scale entre 7-10 geralmente funciona melhor
                - Teste diferentes seeds até encontrar uma que funcione bem
                - O preset "Alta Qualidade" geralmente funciona melhor
                """)
            
            # Mostra exemplos de seeds que funcionam bem
            with gr.Accordion("Seeds que funcionam bem", open=False):
                gr.Markdown("""
                ### Seeds testadas que geram bons resultados:
                
                **Para Paula:**
                - 42689753: Paula em vestido azul
                - 78942561: Paula em ambiente externo
                - 15983264: Close-up de Paula sorrindo
                
                **Para Vivi:**
                - 36798245: Vivi em traje formal
                - 65123987: Vivi em ambiente urbano
                - 93254168: Close-up de Vivi
                
                **Para ambos juntos:**
                - 25874136: Casal em um café
                - 78963214: Passeando em um parque
                - 46125893: Em um restaurante à noite
                """)

    # Define a lógica para mostrar/esconder controles baseado na seleção
    def update_lora_controls(option):
        return {
            lora_scale_1: gr.update(visible=option in ["Paula", "Ambos"]),
            lora_scale_2: gr.update(visible=option in ["Vivi", "Ambos"]),
            lora_balance: gr.update(visible=option == "Ambos")
        }
    
    lora_option.change(
        update_lora_controls,
        inputs=[lora_option],
        outputs=[lora_scale_1, lora_scale_2, lora_balance]
    )

    generate_button.click(
        run_lora,
        inputs=[
            prompt, cfg_scale, steps, randomize_seed, seed, width, height, 
            lora_option, lora_scale_1, lora_scale_2, lora_balance, 
            translate_prompt, use_negative_prompt, quality_preset
        ],
        outputs=[result, seed, final_prompt],
    )

app.queue()
app.launch(share=True)