Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from preprocess import process_dataset
 import subprocess
 import zipfile
+import shutil
 import time
 
 def train_lora_interface(
@@ -11,7 +12,8 @@ def train_lora_interface(
     num_epochs, hub_token, concept_name, description
 ):
     if not dataset_input:
-        return "❌ Por favor, envie um ZIP ou selecione imagens."
+        yield "❌ Por favor, envie um ZIP ou selecione imagens."
+        return
     if not concept_name.strip():
         yield "❌ Por favor, defina um nome para o conceito (ex: brenda)."
         return
@@ -19,70 +21,78 @@ def train_lora_interface(
         yield "❌ Por favor, adicione uma descrição base (ex: mulher, 30 anos, cabelo cacheado)."
         return
 
-    concept_name = concept_name.strip().replace(" ", "_")  # Evita espaços
+    concept_name = concept_name.strip().replace(" ", "_")
     full_description = f"{description.strip()}, {concept_name}"
 
-    yield f"🏷️ Conceito: '{concept_name}' → Descrição: '{full_description}'"
-
-    yield "📁 Preparando dataset..."
+    yield f"🏷️ Treinando conceito: '{concept_name}'"
 
-    # Limpa pasta anterior
+    # Pasta de trabalho
     dataset_dir = "processed_data"
-    if os.path.exists(dataset_dir):
-        for f in os.listdir(dataset_dir):
-            fp = os.path.join(dataset_dir, 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}"
     os.makedirs(dataset_dir, exist_ok=True)
 
+    # Limpa pasta anterior
+    for item in os.listdir(dataset_dir):
+        item_path = os.path.join(dataset_dir, item)
+        try:
+            if os.path.isfile(item_path) or os.path.islink(item_path):
+                os.unlink(item_path)
+            elif os.path.isdir(item_path):
+                shutil.rmtree(item_path)
+        except Exception as e:
+            yield f"⚠️ Erro ao limpar: {e}"
+
+    # --- ETAPA 1: Processar entrada (ZIP ou múltiplas imagens) ---
     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."
+        # dataset_input é uma lista? Pega o primeiro item
+        zip_file = dataset_input[0] if isinstance(dataset_input, list) else dataset_input
+        
+        if not zipfile.is_zipfile(zip_file):
+            yield "❌ Arquivo não é um ZIP válido."
             return
-
-        with zipfile.ZipFile(zip_path, 'r') as z:
+        
+        yield "📦 Descompactando ZIP..."
+        with zipfile.ZipFile(zip_file, 'r') as z:
             z.extractall(dataset_dir)
-        yield f"✅ ZIP descompactado! {len(z.namelist())} arquivos."
+        yield f"✅ ZIP extraído! {len(z.namelist())} arquivos."
 
     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)
-            image_count += 1
-        yield f"✅ {image_count} imagens copiadas."
-
-    # Verifica imagens
-    image_files = [f for f in os.listdir(dataset_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        image_files = dataset_input if isinstance(dataset_input, list) else [dataset_input]
+        yield f"🖼️ Recebidas {len(image_files)} imagens. Copiando..."
+
+        for uploaded_file in image_files:
+            # O Gradio dá um objeto com .name (caminho completo)
+            if hasattr(uploaded_file, 'name'):
+                src_path = uploaded_file.name
+                filename = os.path.basename(src_path)
+                dest_path = os.path.join(dataset_dir, filename)
+                
+                # COPIA (não renomeia) porque /tmp é somente leitura
+                shutil.copy(src_path, dest_path)
+            else:
+                yield f"⚠️ Arquivo inválido: {uploaded_file}"
+        
+        yield f"✅ {len(image_files)} imagens copiadas para {dataset_dir}."
+
+    # --- ETAPA 2: Verifica e gera legendas ---
+    image_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.webp')
+    image_files = [f for f in os.listdir(dataset_dir) if f.lower().endswith(image_extensions)]
+
     if len(image_files) == 0:
-        yield "❌ Nenhuma imagem encontrada. Envie JPG, PNG ou ZIP com imagens."
+        yield "❌ Nenhuma imagem encontrada. Envie imagens válidas (.jpg, .png, etc.)."
         return
 
-    yield f"🖼️ Dataset pronto com {len(image_files)} imagens. Gerando legendas personalizadas..."
+    yield f"📝 Gerando legendas personalizadas para {len(image_files)} imagens..."
 
-    # Gera .txt com base no conceito (sem depender só do BLIP)
+    # Gera .txt com descrição + conceito
     for img_name in image_files:
         txt_path = os.path.join(dataset_dir, os.path.splitext(img_name)[0] + ".txt")
-        if not os.path.exists(txt_path):  # Só cria se não existir
+        if not os.path.exists(txt_path):
             with open(txt_path, "w", encoding="utf-8") as f:
                 f.write(full_description)
 
-    # Atualiza legendas com BLIP (opcional, mas mantém o conceito)
-    try:
-        from preprocess import process_dataset as run_blip
-        run_blip(None, dataset_dir, generate_captions=False)  # Já temos legendas boas!
-        yield "📌 Legendas personalizadas aplicadas com sucesso!"
-    except Exception as e:
-        yield f"⚠️ Falha ao rodar BLIP: {str(e)}. Usando legendas manuais."
-
-    yield "🔥 Iniciando treinamento LoRA..."
+    yield "🔍 Legendas aplicadas com sucesso!"
 
+    # --- ETAPA 3: Treinamento ---
     output_dir = "lora-output"
     os.makedirs(output_dir, exist_ok=True)
 
@@ -103,6 +113,8 @@ def train_lora_interface(
         cmd.append("--hub_model_id")
         cmd.append(f"{concept_name}-lora")
 
+    yield "🔥 Iniciando treinamento LoRA... Isso pode levar minutos."
+
     try:
         process = subprocess.Popen(
             cmd,
@@ -123,20 +135,19 @@ def train_lora_interface(
 
         if process.returncode == 0:
             yield f"""
-🎉 SUCESSO! Modelo LoRA treinado!
+🎉 TREINAMENTO CONCLUÍDO!
 
-🔹 Nome do conceito: **{concept_name}**
-🔹 Use no prompt: `photo of {concept_name} in a garden`
+🔹 Use no prompt: `photo of {concept_name} in the forest`
 🔹 Modelo salvo em: `{output_dir}`
 {'🔹 Publicado no Hub!' if hub_token else ''}
             """
         else:
-            yield f"❌ Treinamento falhou. Código: {process.returncode}\nLogs:\n{log_output[-1000:]}"
+            yield f"❌ Falha no treinamento. Código: {process.returncode}\nLogs:\n{log_output[-1000:]}"
 
     except Exception as e:
-        yield f"💥 Erro ao executar treinamento: {str(e)}"
+        yield f"💥 Erro crítico: {str(e)}"
 
-# Interface atualizada
+# --- Interface Gradio ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🎨 Treinador de LoRA - Hugging Face")
     gr.Markdown("Treine personagens, estilos ou objetos com nome e descrição!")
@@ -151,21 +162,21 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         dataset_input = gr.File(
             label="📤 Envie seu ZIP ou imagens",
-            file_types=[".zip", ".jpg", ".jpeg", ".png"],
+            file_types=[".zip", ".jpg", ".jpeg", ".png", ".bmp", ".webp"],
             file_count="multiple"
         )
 
     gr.Markdown("### 🔖 Identidade do Personagem/Conceito")
     with gr.Row():
         concept_name = gr.Textbox(
-            label="Nome do Conceito (ex: brenda, cyborg_x)",
-            placeholder="Ex: brenda, super_heroi, estilo_anime",
+            label="Nome do Conceito (ex: brenda)",
+            placeholder="Ex: brenda, super_heroi",
             value=""
         )
     with gr.Row():
         description = gr.Textbox(
-            label="Descrição Base (ex: mulher, 30 anos, cabelo cacheado)",
-            placeholder="Ex: young woman, curly hair, green eyes, smiling",
+            label="Descrição Base (ex: woman, curly hair, brown eyes)",
+            placeholder="Ex: young black woman, realistic, warm smile",
             lines=2
         )
 
@@ -194,6 +205,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=output
     )
 
+# Ativa suporte a yield
 demo.queue()
+
 if __name__ == "__main__":
     demo.launch()