Update main.py

import gradio as gr
import os
import shutil
import subprocess

UPLOAD_DIR = "training_images"
OUTPUT_DIR = "lora_output"

def train_lora(images, learning_rate, num_epochs, rank):
if os.path.exists(UPLOAD_DIR):
shutil.rmtree(UPLOAD_DIR)
os.makedirs(UPLOAD_DIR, exist_ok=True)

if os.path.exists(OUTPUT_DIR):
shutil.rmtree(OUTPUT_DIR)
os.makedirs(OUTPUT_DIR, exist_ok=True)

for idx, img in enumerate(images):
img.save(os.path.join(UPLOAD_DIR, f"image_{idx}.png"))

cmd = [
"python", "train_lora.py",
"--images_dir", UPLOAD_DIR,
"--output_dir", OUTPUT_DIR,
"--learning_rate", str(learning_rate),
"--num_epochs", str(num_epochs),
"--rank", str(rank),
]
result = subprocess.run(cmd, capture_output=True, text=True)

output_file = os.path.join(OUTPUT_DIR, "lora.safetensors")
if os.path.exists(output_file):
return f"✅ Treinamento finalizado!\nModelo salvo em: {output_file}\n\nLogs:\n{result.stdout}"
else:
return f"❌ Erro no treinamento:\n{result.stderr}"

# Interface Gradio
with gr.Blocks() as demo:
gr.Markdown("# 🖼️ Criador & Treinador de LoRA")
with gr.Row():
image_input = gr.File(file_types=[".png", ".jpg", ".jpeg"], file_types_display="images", file_count="multiple", label="Envie suas imagens (10–50)")
with gr.Row():
learning_rate = gr.Number(value=1e-4, label="Learning Rate")
num_epochs = gr.Number(value=10, label="Número de Epochs")
rank = gr.Number(value=4, label="Rank do LoRA")
with gr.Row():
train_button = gr.Button("🚀 Treinar LoRA")
output_text = gr.Textbox(label="Saída", lines=15)

train_button.click(fn=train_lora, inputs=[image_input, learning_rate, num_epochs, rank], outputs=output_text)

# 👇 IMPORTANTE: Não usar demo.launch(), apenas expor a variável demo

main.py

@@ -1,56 +0,0 @@
-import gradio as gr
-import os
-import shutil
-import subprocess
-
-UPLOAD_DIR = "training_images"
-OUTPUT_DIR = "lora_output"
-
-def train_lora(images, learning_rate, num_epochs, rank):
-    # Limpa pastas antigas
-    if os.path.exists(UPLOAD_DIR):
-        shutil.rmtree(UPLOAD_DIR)
-    os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-    if os.path.exists(OUTPUT_DIR):
-        shutil.rmtree(OUTPUT_DIR)
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-    # Salva imagens enviadas
-    for idx, img in enumerate(images):
-        img.save(os.path.join(UPLOAD_DIR, f"image_{idx}.png"))
-
-    # Executa o script de treinamento
-    cmd = [
-        "python", "train_lora.py",
-        "--images_dir", UPLOAD_DIR,
-        "--output_dir", OUTPUT_DIR,
-        "--learning_rate", str(learning_rate),
-        "--num_epochs", str(num_epochs),
-        "--rank", str(rank),
-    ]
-    result = subprocess.run(cmd, capture_output=True, text=True)
-
-    # Retorna logs e link para baixar o LoRA
-    output_file = os.path.join(OUTPUT_DIR, "lora.safetensors")
-    if os.path.exists(output_file):
-        return f"✅ Treinamento finalizado!\nModelo salvo em: {output_file}\n\nLogs:\n{result.stdout}"
-    else:
-        return f"❌ Erro no treinamento:\n{result.stderr}"
-
-# Interface Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# 🖼️ Criador & Treinador de LoRA")
-    with gr.Row():
-        image_input = gr.File(file_types=[".png", ".jpg", ".jpeg"], file_types_display="images", file_count="multiple", label="Envie suas imagens (10–50)")
-    with gr.Row():
-        learning_rate = gr.Number(value=1e-4, label="Learning Rate")
-        num_epochs = gr.Number(value=10, label="Número de Epochs")
-        rank = gr.Number(value=4, label="Rank do LoRA")
-    with gr.Row():
-        train_button = gr.Button("🚀 Treinar LoRA")
-    output_text = gr.Textbox(label="Saída", lines=15)
-
-    train_button.click(fn=train_lora, inputs=[image_input, learning_rate, num_epochs, rank], outputs=output_text)
-
-demo.launch()