Update app.py

app.py

@@ -1,9 +1,121 @@
+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
+
+# 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
+lora_models = {
+    "Paula": {
+        "repo": "vcollos/Paula2",
+        "weights": "Paula P.safetensors",
+        "trigger_word": "paula, woman with long blonde hair"  # Trigger word para Paula (mulher)
+    },
+    "Vivi": {
+        "repo": "vcollos/Vivi",
+        "weights": "Vivi.safetensors",
+        "trigger_word": "vivi, man with dark hair"  # Trigger word corrigida para Vivi (homem)
+    }
+}
+
+# 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)
+        print(f"✅ LoRA {name} carregado")
+    except Exception as e:
+        print(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:
+        print(f"❌ Erro no upload da imagem: {e}")
+        return None
+
 @spaces.GPU(duration=80)
-def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, progress=gr.Progress(track_tqdm=True)):
+def translate_text(text, source_lang="pt", target_lang="en"):
+    """
+    Traduz o texto de português para inglês usando a API do LibreTranslate.
+    Se a tradução falhar, retorna o texto original.
+    """
+    try:
+        # Usa a API pública do LibreTranslate
+        url = "https://libretranslate.com/translate"
+        payload = {
+            "q": text,
+            "source": source_lang,
+            "target": target_lang,
+            "format": "text"
+        }
+        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
+            
+    except Exception as e:
+        print(f"❌ Erro ao traduzir texto: {e}")
+        return text
+
+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)):
     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
+    
     # Trunca o prompt para 77 tokens para evitar erro do CLIP
     prompt_tokens = prompt.split()[:77]
     prompt = " ".join(prompt_tokens)
@@ -11,22 +123,34 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     # Define qual LoRA usar com base na seleção do usuário
     selected_loras = []
     adapter_weights = []
-
+    
+    # Modificado para melhorar a mescla de personagens
     if lora_option == "Paula":
         selected_loras.append("Paula")
         adapter_weights.append(lora_scale_1)
+        prompt = f"{lora_models['Paula']['trigger_word']} {prompt}"
+        
     elif lora_option == "Vivi":
         selected_loras.append("Vivi")
         adapter_weights.append(lora_scale_2)
+        prompt = f"{lora_models['Vivi']['trigger_word']} {prompt}"
+        
     elif lora_option == "Ambos":
         selected_loras = ["Paula", "Vivi"]
-        # Ajusta os pesos para garantir equilíbrio
-        total_scale = lora_scale_1 + lora_scale_2
-        adapter_weights = [lora_scale_1 / total_scale, lora_scale_2 / total_scale]
+        adapter_weights = [lora_scale_1, lora_scale_2]
+        
+        # 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 para tentar melhorar a nitidez
+    # Gera a imagem com precisão de 16 bits
     with torch.autocast("cuda"):
         image = pipe(
             prompt=prompt,
@@ -34,7 +158,8 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
             guidance_scale=cfg_scale,
             width=width,
             height=height,
-            generator=generator
+            generator=generator,
+            cross_attention_kwargs=cross_attention_kwargs
         ).images[0]
 
     # Define um nome único para a imagem
@@ -54,13 +179,16 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     # Salva todos os metadados no Supabase
     try:
         response = supabase.table("images").insert({
-            "prompt": prompt,
+            "prompt": original_prompt,  # Salva o prompt original
+            "full_prompt": prompt,      # Salva o prompt completo com trigger words
+            "translated": auto_translate,  # Indica se o prompt foi traduzido
             "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,
             "image_url": image_url,
             "created_at": datetime.utcnow().isoformat()
         }).execute()
@@ -73,4 +201,52 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     except Exception as e:
         print(f"❌ Erro ao salvar metadados no Supabase: {e}")
 
-    return image, seed
+    return image, seed
+
+# 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 (máx 77 caracteres)", lines=3)
+            generate_button = gr.Button("Gerar")
+            
+            with gr.Accordion("Configurações Básicas", open=True):
+                cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
+                steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=32)
+                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(False, 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(["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")
+
+        with gr.Column(scale=2):
+            result = gr.Image(label="Generated Image")
+            gr.Markdown("""
+            ### 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.
+            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
+            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],
+    )
+
+app.queue()
+app.launch(share=True)