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