| | 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}" |
| |
|
| | 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 |
