Update app.py

app.py

@@ -10,7 +10,7 @@ from gradio_client import Client as client_gradio
 from supabase import create_client, Client
 from datetime import datetime
 
-# Inicializa supabase
+# Inicializa Supabase
 url: str = os.getenv('SUPABASE_URL')
 key: str = os.getenv('SUPABASE_KEY')
 supabase: Client = create_client(url, key)
@@ -18,7 +18,7 @@ supabase: Client = create_client(url, key)
 # Obtém token da Hugging Face
 hf_token = os.getenv("HF_TOKEN")
 
-# Inicializa o modelo base
+# 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)
 
@@ -37,9 +37,25 @@ except Exception as e:
 
 # Função para combinar os pesos dos LoRA
 def combine_lora_weights(pipe, weight_1, weight_2):
-    for name, module in pipe.named_modules():  # Percorre os módulos do pipeline
-        if hasattr(module, "weight") and module.weight is not None:
-            module.weight.data = weight_1 * module.weight.data + weight_2 * module.weight.data
+    """
+    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")
 
@@ -68,13 +84,13 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
         )
         raise gr.Error("🚫 Requisição não autorizada!")
 
-    # Atualiza barra de progresso (0% no início)
+    # Atualiza a barra de progresso (0% no início)
     progress(0, "Iniciando a geração de imagem...")
 
-    # Combina os LoRA weights corretamente
+    # Corrige a fusão dos pesos LoRA
     combine_lora_weights(pipe, lora_scale_1, lora_scale_2)
 
-    # Gera imagem com o pipeline
+    # Gera a imagem com o modelo
     image = pipe(
         prompt=f"{prompt}",
         num_inference_steps=steps,
@@ -85,35 +101,7 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
         max_sequence_length=512
     ).images[0]
 
-    # Salva a imagem em um diretório temporário
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    image_filename = f"generated_image_{timestamp}.png"
-    image_path = os.path.join("/tmp/gradio", image_filename)
-
-    # Adiciona metadados à imagem
-    new_metadata_string = f"{prompt}\nNegative prompt: none\nSteps: {steps}, CFG scale: {cfg_scale}, Seed: {seed}, Lora hashes: AndroFlux-v19, VitorCollos"
-    metadata = PngImagePlugin.PngInfo()
-    metadata.add_text("parameters", new_metadata_string)
-
-    # Salva a imagem gerada
-    image.save(image_path, pnginfo=metadata)
-
-    # Registra a imagem no Supabase
-    try:
-        if "girl" not in prompt and "woman" not in prompt:
-            # Salva a imagem no Supabase Storage
-            response = supabase.storage.from_('generated_images').upload(image_filename, image_path, file_options={"content-type": "image/png;charset=UTF-8"})
-            print(response.dict)
-            
-            # Registra a requisição no Supabase
-            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, "image_url": response.full_path})
-            .execute()
-            )
-    except Exception as error:
-        print("⚠️ Erro ao salvar no Supabase:", error)
-
-    yield image, seed
+    return image, seed
 
 # Interface Gradio
 gr_theme = os.getenv("THEME")