# app.py # ⚠️ NÃO REMOVA: Força instalação de pacotes críticos antes de qualquer import import subprocess import sys import os import shutil def install_packages(): print("🔧 Forçando reinstalação de diffusers e huggingface_hub...") try: subprocess.check_call([ sys.executable, "-m", "pip", "install", "--force-reinstall", "diffusers>=0.26.0", "huggingface-hub>=0.20.0", "accelerate", "peft", "torch==2.3.0", "transformers==4.40.0" ]) print("✅ Pacotes essenciais instalados!") except Exception as e: print(f"❌ Falha na instalação forçada: {e}") # Executa apenas uma vez por sessão if not os.path.exists("/tmp/packages_installed"): install_packages() with open("/tmp/packages_installed", "w") as f: f.write("ok") # Agora sim, imports seguros import gradio as gr from preprocess import process_dataset import subprocess import zipfile import time def train_lora_interface( dataset_input, input_type, model_name, lora_rank, learning_rate, num_epochs, hub_token, concept_name, description ): if not dataset_input: 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 if not description.strip(): yield "❌ Por favor, adicione uma descrição base (ex: mulher, 30 anos, cabelo cacheado)." return concept_name = concept_name.strip().replace(" ", "_") full_description = f"{description.strip()}, {concept_name}" yield f"🏷️ Treinando conceito: '{concept_name}' → Prompt: [photo of {concept_name}]" # Pasta de trabalho dataset_dir = "processed_data" 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 --- if input_type == "Upload de ZIP": 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 yield "📦 Descompactando ZIP..." with zipfile.ZipFile(zip_file, 'r') as z: z.extractall(dataset_dir) yield f"✅ ZIP extraído! {len(z.namelist())} arquivos." else: # Múltiplas imagens 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: if hasattr(uploaded_file, 'name'): src_path = uploaded_file.name filename = os.path.basename(src_path) dest_path = os.path.join(dataset_dir, filename) shutil.copy(src_path, dest_path) # Usa copy, não rename else: yield f"⚠️ Arquivo inválido: {uploaded_file}" yield f"✅ {len(image_files)} imagens copiadas." # --- ETAPA 2: Verifica imagens 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 arquivos válidos." return yield f"📝 Aplicando legenda base: '{full_description}'" 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): with open(txt_path, "w", encoding="utf-8") as f: f.write(full_description) yield "🔍 Legendas aplicadas com sucesso!" # --- ETAPA 3: 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(f"{concept_name}-lora") yield "🔥 Iniciando treinamento LoRA... Isso pode levar alguns minutos." try: process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, bufsize=1, encoding='utf-8' ) log_output = "" for line in process.stdout: log_output += line 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! 🔹 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:]}" except Exception as e: yield f"💥 Erro ao executar treinamento: {str(e)}" # --- Interface Gradio --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# 🎨 Treinador de LoRA - Hugging Face") gr.Markdown("Treine seu próprio modelo com nome, descrição e imagens!") with gr.Row(): 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", ".bmp", ".webp"], file_count="multiple" ) gr.Markdown("### 🔖 Identidade do Personagem/Conceito") with gr.Row(): concept_name = gr.Textbox( label="Nome do Conceito (ex: brenda)", placeholder="Ex: brenda, cyborg_x, estilo_pintura", value="" ) with gr.Row(): description = gr.Textbox( label="Descrição Base (ex: woman, curly hair, realistic)", placeholder="Ex: young black woman, warm smile, detailed face", lines=2 ) gr.Markdown("### ⚙️ Configurações do Treinamento") with gr.Row(): model_name = gr.Dropdown( ["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="LoRA Rank") learning_rate = gr.Number(value=1e-4, label="Taxa de Aprendizado") 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="📦 Logs e Resultado", lines=12) btn.click( train_lora_interface, inputs=[ dataset_input, input_type, model_name, lora_rank, learning_rate, num_epochs, hub_token, concept_name, description ], outputs=output ) # Ativa suporte a yield demo.queue() if __name__ == "__main__": demo.launch()