Update app.py

app.py

@@ -3,43 +3,94 @@ import os
 import gradio as gr
 from preprocess import process_dataset
 import subprocess
+import zipfile
 import time
 
-def train_lora_interface(dataset_zip, model_name, lora_rank, learning_rate, num_epochs, hub_token):
-    if not dataset_zip:
-        return "❌ Por favor, envie um ZIP com suas imagens."
-
-    # Etapa 1: Pré-processamento
-    yield "🔄 Descompactando e processando dataset..."
-    try:
-        dataset_dir = process_dataset(dataset_zip, "processed_data")
-        image_count = len([f for f in os.listdir(dataset_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
-        yield f"✅ Dataset processado: {image_count} imagens encontradas. Iniciando treinamento..."
-
-        # Etapa 2: Configura treinamento
-        output_dir = "lora-output"
-        os.makedirs(output_dir, exist_ok=True)
-
-        cmd = [
-            "python", "train_lora.py",
-            "--dataset_dir", dataset_dir,
-            "--model_name", model_name,
-            "--lora_rank", str(lora_rank),
-            "--learning_rate", str(learning_rate),
-            "--num_epochs", str(num_epochs),
-            "--batch_size", "1",
-            "--output_dir", output_dir
-        ]
-
-        if hub_token:
-            os.environ["HF_TOKEN"] = hub_token
-            cmd.append("--push_to_hub")
-            cmd.append("--hub_model_id")
-            cmd.append("my-lora-model")
-
-        # Etapa 3: Executa treinamento
-        yield "🔥 Treinando modelo... Isso pode levar alguns minutos."
+def train_lora_interface(dataset_input, input_type, model_name, lora_rank, learning_rate, num_epochs, hub_token):
+    if not dataset_input:
+        return "❌ Por favor, envie um ZIP ou selecione imagens."
+
+    yield "📁 Preparando dataset..."
+
+    # Pasta temporária
+    os.makedirs("processed_data", exist_ok=True)
+    for f in os.listdir("processed_data"):
+        fp = os.path.join("processed_data", f)
+        try:
+            if os.path.isfile(fp) or os.path.islink(fp):
+                os.unlink(fp)
+            elif os.path.isdir(fp):
+                os.rmtree(fp)
+        except Exception as e:
+            yield f"⚠️ Erro ao limpar: {e}"
+
+    dataset_dir = "processed_data"
+
+    if input_type == "Upload de ZIP":
+        zip_path = dataset_input
+        if not zipfile.is_zipfile(zip_path):
+            yield "❌ Arquivo enviado não é um ZIP válido."
+            return
         
+        # Descompacta
+        with zipfile.ZipFile(zip_path, 'r') as z:
+            z.extractall(dataset_dir)
+        yield f"✅ ZIP descompactado! {len(z.namelist())} arquivos extraídos."
+
+    else:  # Múltiplas imagens
+        image_count = 0
+        for img_path in dataset_input:
+            dest_path = os.path.join(dataset_dir, os.path.basename(img_path.name))
+            os.rename(img_path.name, dest_path)  # Move para processed_data
+            
+            # Cria .txt vazio (BLIP vai preencher depois)
+            txt_path = os.path.splitext(dest_path)[0] + ".txt"
+            if not os.path.exists(txt_path):
+                with open(txt_path, "w") as f:
+                    f.write("person")
+            image_count += 1
+        yield f"✅ {image_count} imagens copiadas para o dataset."
+
+    # Conta imagens processadas
+    image_files = [f for f in os.listdir(dataset_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    if len(image_files) == 0:
+        yield "❌ Nenhuma imagem encontrada. Envie JPG, PNG ou ZIP com imagens."
+        return
+
+    yield f"🖼️ Dataset pronto com {len(image_files)} imagens. Iniciando pré-processamento..."
+
+    # Gera captions com BLIP (reutiliza preprocess.py)
+    try:
+        from preprocess import process_dataset as run_blip
+        run_blip(None, dataset_dir, generate_captions=True)  # Já temos as imagens, só gera captions
+    except Exception as e:
+        yield f"⚠️ Falha ao gerar legendas: {str(e)}. Continuando com captions existentes."
+
+    yield "🔥 Iniciando treinamento LoRA..."
+
+    # Comando de treinamento
+    output_dir = "lora-output"
+    os.makedirs(output_dir, exist_ok=True)
+
+    cmd = [
+        "python", "train_lora.py",
+        "--dataset_dir", dataset_dir,
+        "--model_name", model_name,
+        "--lora_rank", str(lora_rank),
+        "--learning_rate", str(learning_rate),
+        "--num_epochs", str(num_epochs),
+        "--batch_size", "1",
+        "--output_dir", output_dir
+    ]
+
+    if hub_token:
+        os.environ["HF_TOKEN"] = hub_token
+        cmd.append("--push_to_hub")
+        cmd.append("--hub_model_id")
+        cmd.append("my-lora-model")
+
+    # Executa treinamento
+    try:
         process = subprocess.Popen(
             cmd,
             stdout=subprocess.PIPE,
@@ -52,28 +103,39 @@ def train_lora_interface(dataset_zip, model_name, lora_rank, learning_rate, num_
         log_output = ""
         for line in process.stdout:
             log_output += line
-            # Mostra os últimos logs a cada 50 linhas
-            if len(log_output.split('\n')) % 20 == 0:
-                yield f"📝 Treinando...\n{log_output[-500:]}"
+            if "loss" in line.lower() or "epoch" in line.lower():
+                yield f"📊 {line.strip()}"
 
         process.wait()
 
         if process.returncode == 0:
-            yield f"🎉 Treinamento concluído! Modelo salvo em `{output_dir}`"
+            yield f"🎉 SUCESSO! Modelo LoRA treinado e salvo em `{output_dir}`"
         else:
-            yield f"❌ Falha no treinamento. Código: {process.returncode}\nÚltimos logs:\n{log_output[-1000:]}"
-    
+            yield f"❌ Treinamento falhou com código {process.returncode}.\nLogs:\n{log_output[-1000:]}"
+
     except Exception as e:
-        yield f"💥 Erro crítico: {str(e)}\n\nVerifique seus arquivos e tente novamente."
+        yield f"💥 Erro ao executar treinamento: {str(e)}"
 
-# Interface
+# Interface com opção de entrada
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🎨 Treinador de LoRA - Hugging Face")
-    gr.Markdown("Treine modelos personalizados com apenas algumas imagens!")
+    gr.Markdown("Envie suas imagens via ZIP ou múltiplos arquivos!")
 
     with gr.Row():
-        dataset_zip = gr.File(label="📤 Envie um ZIP com suas imagens (.jpg, .png)")
+        input_type = gr.Radio(
+            ["Upload de ZIP", "Selecionar várias imagens"],
+            label="Tipo de Entrada",
+            value="Upload de ZIP"
+        )
+
+    with gr.Row():
+        dataset_input = gr.File(
+            label="📤 Envie seu ZIP ou imagens",
+            file_types=[".zip", ".jpg", ".jpeg", ".png"],
+            file_count="multiple"  # Permite múltiplos arquivos!
+        )
 
+    gr.Markdown("### ⚙️ Configurações do Treinamento")
     with gr.Row():
         model_name = gr.Dropdown(
             ["runwayml/stable-diffusion-v1-5"],
@@ -87,15 +149,15 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     hub_token = gr.Textbox(label="🔐 Token do Hugging Face (opcional)", type="password")
 
     btn = gr.Button("🚀 Iniciar Treinamento", variant="primary")
-    output = gr.Textbox(label="📊 Status / Logs", lines=10)
+    output = gr.Textbox(label="📦 Logs e Status", lines=12)
 
     btn.click(
         train_lora_interface,
-        inputs=[dataset_zip, model_name, lora_rank, learning_rate, num_epochs, hub_token],
+        inputs=[dataset_input, input_type, model_name, lora_rank, learning_rate, num_epochs, hub_token],
         outputs=output
     )
 
-# Ativa fila para suporte a `yield`
+# Ativa fila para yield
 demo.queue()
 
 if __name__ == "__main__":