Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os
 import gradio as gr
 from preprocess import process_dataset
@@ -5,18 +6,20 @@ import subprocess
 import time
 
 def train_lora_interface(dataset_zip, model_name, lora_rank, learning_rate, num_epochs, hub_token):
-    # 1. Pré-processamento
-    with gr.Progress() as progress:
-        progress(0, "Descompactando e processando dataset...")
+    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")
-        
-        # 2. Configura treinamento
-        progress(0.3, "Configurando treinamento LoRA...")
+        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)
-        
-        # 3. Executa treinamento
-        progress(0.5, "Treinando modelo (isso pode levar horas)...")
+
         cmd = [
             "python", "train_lora.py",
             "--dataset_dir", dataset_dir,
@@ -24,66 +27,76 @@ def train_lora_interface(dataset_zip, model_name, lora_rank, learning_rate, num_
             "--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."
         
         process = subprocess.Popen(
-            cmd, 
-            stdout=subprocess.PIPE, 
+            cmd,
+            stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,
-            universal_newlines=True
+            universal_newlines=True,
+            bufsize=1,
+            encoding='utf-8'
         )
-        
-        logs = ""
+
+        log_output = ""
         for line in process.stdout:
-            logs += line
-            progress(0.7, f"Treinando...\n{logs[-200:]}")
-        
-        # 4. Finalização
-        progress(0.9, "Subindo para Hugging Face Hub...")
-        if hub_token:
-            from huggingface_hub import upload_folder
-            upload_folder(
-                repo_id="my-lora-model",
-                folder_path=output_dir,
-                token=hub_token
-            )
-        
-        progress(1.0, "Treinamento concluído com sucesso!")
-        return f"Modelo salvo em: {output_dir}\nLogs: {logs[-500:]}"
+            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:]}"
+
+        process.wait()
 
-# Interface Gradio
-with gr.Blocks(title="LoRA Trainer - Hugging Face") as demo:
-    gr.Markdown("# 🚀 Treinador de LoRA para Stable Diffusion")
-    gr.Markdown("Treine seus próprios modelos LoRA diretamente no Hugging Face Spaces!")
+        if process.returncode == 0:
+            yield f"🎉 Treinamento concluído! Modelo salvo em `{output_dir}`"
+        else:
+            yield f"❌ Falha no treinamento. Código: {process.returncode}\nÚltimos logs:\n{log_output[-1000:]}"
     
-    with gr.Tab("Configuração"):
-        dataset_zip = gr.File(label="Dataset (ZIP com imagens)", file_types=['.zip'])
+    except Exception as e:
+        yield f"💥 Erro crítico: {str(e)}\n\nVerifique seus arquivos e tente novamente."
+
+# Interface
+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!")
+
+    with gr.Row():
+        dataset_zip = gr.File(label="📤 Envie um ZIP com suas imagens (.jpg, .png)")
+
+    with gr.Row():
         model_name = gr.Dropdown(
-            ["runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1"], 
+            ["runwayml/stable-diffusion-v1-5"],
             value="runwayml/stable-diffusion-v1-5",
             label="Modelo Base"
         )
-        lora_rank = gr.Slider(4, 64, value=4, step=4, label="Rank LoRA")
+        lora_rank = gr.Slider(4, 64, value=4, step=4, label="LoRA Rank")
         learning_rate = gr.Number(value=1e-4, label="Taxa de Aprendizado")
-        num_epochs = gr.Slider(1, 50, value=10, step=1, label="Épocas")
-        hub_token = gr.Textbox(label="Token Hugging Face (opcional)", type="password")
-    
-    with gr.Tab("Treinamento"):
-        start_btn = gr.Button("🚀 Iniciar Treinamento", variant="primary")
-        output = gr.Textbox(label="Logs do Treinamento")
-    
-    start_btn.click(
-        fn=train_lora_interface,
+        num_epochs = gr.Slider(1, 30, value=10, step=1, label="Épocas")
+
+    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)
+
+    btn.click(
+        train_lora_interface,
         inputs=[dataset_zip, model_name, lora_rank, learning_rate, num_epochs, hub_token],
         outputs=output
     )
 
+# Ativa fila para suporte a `yield`
+demo.queue()
+
 if __name__ == "__main__":
     demo.launch()