Update app.py

app.py

@@ -1,7 +1,7 @@
 import spaces
 import gradio as gr
 import torch
-from PIL import Image, PngImagePlugin
+from PIL import Image
 from diffusers import DiffusionPipeline
 import random
 import os
@@ -20,45 +20,25 @@ 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.bfloat16)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16, use_safetensors=True)
 
-# Carrega os adaptadores LoRA
+# Move o modelo para GPU
+pipe.to("cuda")
+
+# Carrega os adaptadores LoRA corretamente
 lora_repo_1 = "markury/AndroFlux"
 lora_repo_2 = "vcollos/VitorCollos"
 
 try:
-    pipe.load_lora_weights(lora_repo_1, weight_name="AndroFlux-v19.safetensors")
+    pipe.load_lora_weights(lora_repo_1, weight_name="AndroFlux-v19.safetensors", adapter_name="AndroFlux-v19")
     print("✅ Primeiro LoRA carregado")
 
-    pipe.load_lora_weights(lora_repo_2, weight_name="Vitor.safetensors")
+    pipe.load_lora_weights(lora_repo_2, weight_name="Vitor.safetensors", adapter_name="Vitor")
     print("✅ Segundo LoRA carregado")
 except Exception as e:
     print(f"❌ Erro ao carregar os LoRA adapters: {e}")
 
-# Função para combinar os pesos dos LoRA
-def combine_lora_weights(pipe, weight_1, weight_2):
-    """
-    Combina os pesos de dois LoRA adapters sem depender de named_modules().
-    """
-    # Verifica quais componentes do pipeline têm os LoRA adapters
-    lora_layers = []
-    
-    if hasattr(pipe, "text_encoder"):
-        lora_layers.append(pipe.text_encoder)
-    if hasattr(pipe, "vae"):
-        lora_layers.append(pipe.vae)
-    if hasattr(pipe, "transformer"):
-        lora_layers.append(pipe.transformer)
-    
-    # Aplica a fusão dos pesos LoRA apenas nos componentes relevantes
-    for module in lora_layers:
-        for param in module.parameters():
-            param.data = weight_1 * param.data + weight_2 * param.data
-
-    print("✅ Pesos LoRA combinados com sucesso!")
-
-pipe.to("cuda")
-
+# Define seed máximo
 MAX_SEED = 2**32 - 1
 
 @spaces.GPU(duration=80)
@@ -68,7 +48,7 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
 
-    # Moderação de texto
+    # Moderação de texto (evita prompts inadequados)
     moderation_client = client_gradio("duchaba/Friendly_Text_Moderation")
     result = moderation_client.predict(
         msg=f"{prompt}",
@@ -78,27 +58,32 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
 
     if float(json.loads(result[1])['sexual_minors']) > 0.03:
         print('🔴 Conteúdo não permitido')
-        response_data = (supabase.table("requests")
-        .insert({"prompt": prompt, "cfg_scale": cfg_scale, "steps": steps, "randomized_seed": randomize_seed, "seed": seed, "lora_scale_1": lora_scale_1, "lora_scale_2": lora_scale_2, "moderated": 'true'})
-        .execute()
-        )
+        supabase.table("requests").insert({
+            "prompt": prompt,
+            "cfg_scale": cfg_scale,
+            "steps": steps,
+            "randomized_seed": randomize_seed,
+            "seed": seed,
+            "lora_scale_1": lora_scale_1,
+            "lora_scale_2": lora_scale_2,
+            "moderated": 'true'
+        }).execute()
         raise gr.Error("🚫 Requisição não autorizada!")
 
-    # Atualiza a barra de progresso (0% no início)
+    # Atualiza a barra de progresso
     progress(0, "Iniciando a geração de imagem...")
 
-    # Corrige a fusão dos pesos LoRA
-    combine_lora_weights(pipe, lora_scale_1, lora_scale_2)
+    # Aplica os adaptadores LoRA corretamente
+    pipe.set_adapters(["AndroFlux-v19", "Vitor"], adapter_weights=[lora_scale_1, lora_scale_2])
 
     # Gera a imagem com o modelo
     image = pipe(
-        prompt=f"{prompt}",
+        prompt=prompt,
         num_inference_steps=steps,
         guidance_scale=cfg_scale,
         width=width,
         height=height,
-        generator=generator,
-        max_sequence_length=512
+        generator=generator
     ).images[0]
 
     return image, seed
@@ -113,8 +98,8 @@ with gr.Blocks(theme=gr_theme) as app:
             generate_button = gr.Button("Gerar")
             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=25)
-            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=896)
-            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1152)
+            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=768)  # Reduzido para evitar falta de VRAM
+            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=768)  # Reduzido para evitar falta de VRAM
             randomize_seed = gr.Checkbox(False, label="Randomize seed")
             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=556215326)
             lora_scale_1 = gr.Slider(label="LoRA Scale (AndroFlux)", minimum=0, maximum=1, step=0.01, value=0.7)