Update app.py

app.py

@@ -12,6 +12,13 @@ from gradio_client import Client as client_gradio
 from supabase import create_client, Client
 from datetime import datetime
 import requests
+import logging
+from google.cloud import translate_v2 as translate
+import google.generativeai as genai
+
+# 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')
@@ -21,6 +28,25 @@ supabase: Client = create_client(url, key)
 # Obtém token da Hugging Face
 hf_token = os.getenv("HF_TOKEN")
 
+# Configura as chaves de API da OpenAI e Google AI
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
+GOOGLE_AI_KEY = os.getenv("GOOGLE_AI_KEY", "")
+
+# Se as chaves não estiverem definidas, usa as fornecidas
+if not OPENAI_API_KEY:
+    OPENAI_API_KEY = "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"
+if not GOOGLE_AI_KEY:
+    GOOGLE_AI_KEY = "AIzaSyD4jxsvKv353nLwVlOYXQBL7Im3-0Vc_6g"
+
+# Configura o Google AI Gemini
+try:
+    genai.configure(api_key=GOOGLE_AI_KEY)
+    google_ai_available = True
+    logger.info("✅ Google AI Gemini configurado com sucesso")
+except Exception as e:
+    google_ai_available = False
+    logger.error(f"❌ Erro ao configurar Google AI Gemini: {e}")
+
 # Inicializa o modelo base FLUX.1-dev
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(
@@ -50,9 +76,9 @@ lora_models = {
 for name, details in lora_models.items():
     try:
         pipe.load_lora_weights(details["repo"], weight_name=details["weights"], adapter_name=name)
-        print(f"✅ LoRA {name} carregado")
+        logger.info(f"✅ LoRA {name} carregado")
     except Exception as e:
-        print(f"❌ Erro ao carregar o LoRA {name}: {e}")
+        logger.error(f"❌ Erro ao carregar o LoRA {name}: {e}")
 
 # Define seed máximo
 MAX_SEED = 2**32 - 1
@@ -73,48 +99,151 @@ def upload_image_to_supabase(image, filename):
         base_url = f"{url}/storage/v1/object/public/images"
         return f"{base_url}/{storage_path}"
     except Exception as e:
-        print(f"❌ Erro no upload da imagem: {e}")
+        logger.error(f"❌ Erro no upload da imagem: {e}")
         return None
 
-def translate_text(text, source_lang="pt", target_lang="en"):
+def optimize_with_openai(text, character_option):
     """
-    Traduz o texto de português para inglês usando a API do LibreTranslate.
-    Se a tradução falhar, retorna o texto original.
+    Otimiza o prompt usando a API da OpenAI.
     """
     try:
-        # Usa a API pública do LibreTranslate
-        url = "https://libretranslate.com/translate"
-        payload = {
-            "q": text,
-            "source": source_lang,
-            "target": target_lang,
-            "format": "text"
+        # Determina o prompt de sistema com base na seleção de personagem
+        if character_option == "Paula":
+            system_prompt = """Você é um especialista em criar prompts para o modelo de imagem FLUX.1-dev. 
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de uma mulher loira chamada Paula. Seu prompt deve ser detalhado, claro e 
+            enfatizar características de uma mulher loira. 
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        elif character_option == "Vivi":
+            system_prompt = """Você é um especialista em criar prompts para o modelo de imagem FLUX.1-dev.
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de um homem de cabelo escuro chamado Vivi. Seu prompt deve ser detalhado, claro e
+            enfatizar características de um homem moreno.
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        else:  # both
+            system_prompt = """Você é um especialista em criar prompts para o modelo de imagem FLUX.1-dev.
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de duas pessoas juntas: uma mulher loira chamada Paula e um homem de cabelo escuro chamado Vivi.
+            Seu prompt deve ser detalhado, claro e deve enfatizar a presença dos DOIS personagens na mesma cena.
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        
+        # Configuração do cabeçalho e corpo da requisição
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {OPENAI_API_KEY}"
         }
-        response = requests.post(url, json=payload)
         
-        if response.status_code == 200:
-            return response.json()["translatedText"]
-        else:
-            print(f"❌ Erro na tradução (código {response.status_code}): {response.text}")
-            return text
-            
+        data = {
+            "model": "gpt-4",
+            "messages": [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": text}
+            ],
+            "temperature": 0.7,
+            "max_tokens": 300
+        }
+        
+        # Faz a requisição para a API da OpenAI
+        response = requests.post(
+            "https://api.openai.com/v1/chat/completions",
+            headers=headers, 
+            json=data
+        )
+        
+        response.raise_for_status()
+        
+        # Extrai e retorna o prompt otimizado
+        optimized_prompt = response.json()["choices"][0]["message"]["content"].strip()
+        logger.info(f"✅ Prompt otimizado via OpenAI: {optimized_prompt}")
+        
+        return optimized_prompt
+    
     except Exception as e:
-        print(f"❌ Erro ao traduzir texto: {e}")
+        logger.error(f"❌ Erro ao otimizar prompt com OpenAI: {e}")
         return text
 
+def optimize_with_google_ai(text, character_option):
+    """
+    Otimiza o prompt usando o Google AI Gemini.
+    """
+    try:
+        # Determina o prompt com base na seleção de personagem
+        if character_option == "Paula":
+            instruction = """Você é um especialista em criar prompts para modelos de imagem AI. 
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de uma mulher loira chamada Paula. Seu prompt deve ser detalhado, claro e 
+            enfatizar características de uma mulher loira. 
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        elif character_option == "Vivi":
+            instruction = """Você é um especialista em criar prompts para modelos de imagem AI.
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de um homem de cabelo escuro chamado Vivi. Seu prompt deve ser detalhado, claro e
+            enfatizar características de um homem moreno.
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        else:  # both
+            instruction = """Você é um especialista em criar prompts para modelos de imagem AI.
+            Traduza e otimize este prompt em português para um prompt em inglês que irá gerar
+            uma imagem de duas pessoas juntas: uma mulher loira chamada Paula e um homem de cabelo escuro chamado Vivi.
+            Seu prompt deve ser detalhado, claro e deve enfatizar a presença dos DOIS personagens na mesma cena.
+            Responda APENAS com o prompt otimizado em inglês, sem explicações ou introduções."""
+        
+        # Configura o modelo Gemini
+        generation_config = {
+            "temperature": 0.7,
+            "top_p": 1,
+            "top_k": 32,
+            "max_output_tokens": 300,
+        }
+        
+        model = genai.GenerativeModel(
+            model_name="gemini-pro",
+            generation_config=generation_config
+        )
+        
+        # Cria a conversa
+        convo = model.start_chat(history=[])
+        convo.send_message(f"{instruction}\n\nPrompt original: {text}")
+        
+        # Obtém a resposta
+        response = convo.last.text
+        logger.info(f"✅ Prompt otimizado via Google AI: {response}")
+        
+        return response.strip()
+    
+    except Exception as e:
+        logger.error(f"❌ Erro ao otimizar prompt com Google AI: {e}")
+        # Tenta o OpenAI como fallback
+        logger.info("🔄 Tentando OpenAI como fallback...")
+        return optimize_with_openai(text, character_option)
+
 @spaces.GPU(duration=80)
-def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, cross_attention_scale, auto_translate, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, lora_balance, ai_option, 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 habilitada
-    if auto_translate:
-        translated_prompt = translate_text(prompt)
-        print(f"📝 Prompt traduzido: {translated_prompt}")
-        prompt = translated_prompt
+    # Determina qual personagem está sendo usado para a otimização
+    character_option = "Paula" if lora_option == "Paula" else "Vivi" if lora_option == "Vivi" else "both"
+    
+    # Processa o prompt de acordo com a opção selecionada
+    processed_prompt = prompt  # Valor padrão
+    
+    if ai_option == "OpenAI":
+        processed_prompt = optimize_with_openai(prompt, character_option)
+    elif ai_option == "Google AI":
+        if google_ai_available:
+            processed_prompt = optimize_with_google_ai(prompt, character_option)
+        else:
+            logger.warning("Google AI indisponível, usando OpenAI como fallback")
+            processed_prompt = optimize_with_openai(prompt, character_option)
+    elif ai_option == "Nenhum":
+        # Mantém o prompt original, sem processamento
+        processed_prompt = prompt
+    
+    # Atualiza o prompt com o texto processado
+    prompt = processed_prompt
     
     # Trunca o prompt para 77 tokens para evitar erro do CLIP
     prompt_tokens = prompt.split()[:77]
@@ -136,19 +265,18 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
         prompt = f"{lora_models['Vivi']['trigger_word']} {prompt}"
         
     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 = [lora_scale_1, lora_scale_2]
+        adapter_weights = [p_weight, v_weight]
         
         # Quando usando ambos, adiciona trigger words específicas para cada personagem
         # e ajusta o prompt para indicar que ambos devem aparecer na mesma cena
         prompt = f"{lora_models['Paula']['trigger_word']} and {lora_models['Vivi']['trigger_word']} together, a woman and a man, {prompt}"
 
     pipe.set_adapters(selected_loras, adapter_weights)
-    
-    # Ajusta a cross-attention quando usando ambos os LoRAs
-    cross_attention_kwargs = {}
-    if lora_option == "Ambos" and cross_attention_scale != 1.0:
-        cross_attention_kwargs = {"scale": cross_attention_scale}
 
     # Gera a imagem com precisão de 16 bits
     with torch.autocast("cuda"):
@@ -158,8 +286,7 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
             guidance_scale=cfg_scale,
             width=width,
             height=height,
-            generator=generator,
-            cross_attention_kwargs=cross_attention_kwargs
+            generator=generator
         ).images[0]
 
     # Define um nome único para a imagem
@@ -168,40 +295,41 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     try:
         image_url = upload_image_to_supabase(image, filename)
         if image_url:
-            print(f"✅ Imagem salva no Supabase: {image_url}")
+            logger.info(f"✅ Imagem salva no Supabase: {image_url}")
         else:
-            print("❌ Erro: URL da imagem retornou None")
-            return image, seed
+            logger.error("❌ Erro: URL da imagem retornou None")
+            return image, seed, prompt
     except Exception as e:
-        print(f"❌ Erro ao fazer upload da imagem: {e}")
-        return image, seed
+        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
+            "processed_prompt": processed_prompt,  # Salva o prompt processado pela IA
             "full_prompt": prompt,      # Salva o prompt completo com trigger words
-            "translated": auto_translate,  # Indica se o prompt foi traduzido
+            "ai_option": ai_option,
             "cfg_scale": cfg_scale,
             "steps": steps,
             "seed": seed,
             "lora_option": lora_option,
             "lora_scale_1": lora_scale_1,
             "lora_scale_2": lora_scale_2,
-            "cross_attention_scale": cross_attention_scale,
+            "lora_balance": lora_balance,
             "image_url": image_url,
             "created_at": datetime.utcnow().isoformat()
         }).execute()
         
         if response.data:
-            print("✅ Metadados salvos no Supabase")
+            logger.info("✅ Metadados salvos no Supabase")
         else:
-            print("❌ Erro: Resposta vazia do Supabase")
+            logger.error("❌ Erro: Resposta vazia do Supabase")
         
     except Exception as e:
-        print(f"❌ Erro ao salvar metadados no Supabase: {e}")
+        logger.error(f"❌ Erro ao salvar metadados no Supabase: {e}")
 
-    return image, seed
+    return image, seed, prompt
 
 # Interface Gradio
 gr_theme = os.getenv("THEME")
@@ -210,7 +338,7 @@ with gr.Blocks(theme=gr_theme) as app:
     
     with gr.Row():
         with gr.Column(scale=2):
-            prompt = gr.TextArea(label="Prompt", placeholder="Digite um prompt (máx 77 caracteres)", lines=3)
+            prompt = gr.TextArea(label="Prompt em Português", placeholder="Digite um prompt em português descrevendo o que você quer", lines=3)
             generate_button = gr.Button("Gerar")
             
             with gr.Accordion("Configurações Básicas", open=True):
@@ -225,27 +353,39 @@ with gr.Blocks(theme=gr_theme) as app:
                 lora_option = gr.Radio(["Nenhum", "Paula", "Vivi", "Ambos"], label="Escolha o LoRA", value="Ambos")
                 lora_scale_1 = gr.Slider(label="LoRA Scale (Paula)", minimum=0, maximum=1, step=0.01, value=0.8)
                 lora_scale_2 = gr.Slider(label="LoRA Scale (Vivi)", minimum=0, maximum=1, step=0.01, value=0.8)
-                cross_attention_scale = gr.Slider(label="Cross-Attention Scale", minimum=0.1, maximum=1.5, step=0.05, value=1.0, 
-                                                 info="Ajuste para controlar o balanceamento dos LoRAs (útil no modo 'Ambos')")
-                auto_translate = gr.Checkbox(True, label="Traduzir automaticamente para inglês", 
-                                           info="Traduz o prompt do português para o inglês antes de gerar a imagem")
+                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.Accordion("Processamento de Prompt", open=True):
+                ai_option = gr.Radio(
+                    ["Nenhum", "OpenAI", "Google AI"], 
+                    label="Método de Otimização", 
+                    value="Google AI" if google_ai_available else "OpenAI"
+                )
 
         with gr.Column(scale=2):
             result = gr.Image(label="Generated Image")
+            final_prompt = gr.Textbox(label="Prompt Final (usado para gerar a imagem)", lines=3)
             gr.Markdown("""
+            ### Como funciona:
+            1. Escreva seu prompt em português
+            2. Escolha entre Google AI (recomendado) ou OpenAI para traduzir e otimizar seu prompt
+            3. O sistema irá transformar sua descrição em português em um prompt profissional em inglês
+            4. O prompt será específico para o(s) personagem(ns) selecionado(s)
+            
             ### Dicas para gerar imagens com ambos personagens:
             1. Use o modo "Ambos" com valores balanceados (ex: 0.8 para ambos os LoRAs)
-            2. Inclua no prompt termos como "um homem e uma mulher juntos", "duas pessoas", etc.
+            2. Inclua no prompt termos como "um homem e uma mulher juntos", "um casal", etc.
             3. Lembre-se: Paula é uma mulher loira e Vivi é um homem de cabelo escuro
-            4. Seu prompt será automaticamente traduzido para inglês (se a opção estiver ativada)
-            5. Ajuste o Cross-Attention Scale para valores como 0.9 ou 1.1 para equilibrar a influência
+            4. Ajuste o Balanço para favorecer um personagem ou outro se necessário
+            5. A otimização via IA produzirá melhores resultados do que prompts diretos
             6. Experimente com seeds diferentes até encontrar uma que funcione bem
             """)
 
     generate_button.click(
         run_lora,
-        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, cross_attention_scale, auto_translate],
-        outputs=[result, seed],
+        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, lora_balance, ai_option],
+        outputs=[result, seed, final_prompt],
     )
 
 app.queue()