Update app.py

app.py

@@ -1,55 +1,150 @@
-import gradio as gr
 import os
-from some_module import train_lora, process_images  # ajuste de acordo com seus imports reais
-
-# Função principal de treinamento
-def start_training(images, model_name, trigger_word, epochs, batch_size, learning_rate, lora_r, lora_alpha, lora_dropout):
-    # processa imagens
-    processed = process_images(images)
-    
-    # chama função de treinamento LoRA
-    result_path = train_lora(
-        images=processed,
-        model_name=model_name,
-        trigger_word=trigger_word,
-        epochs=epochs,
-        batch_size=batch_size,
-        learning_rate=learning_rate,
-        lora_r=lora_r,
-        lora_alpha=lora_alpha,
-        lora_dropout=lora_dropout
-    )
-    return f"Treinamento concluído! Modelo salvo em: {result_path}"
-
-# Criação da interface Gradio com Blocks
+import json
+import uuid
+import threading
+import time
+from datetime import datetime
+from pathlib import Path
+import zipfile
+
+import gradio as gr
+from PIL import Image
+import torch
+
+# --- Classe LoRA Trainer ---
+class LoRAImageTrainer:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.training_jobs = {}
+        self.models_cache = {}
+
+    def get_available_models(self):
+        return [
+            "runwayml/stable-diffusion-v1-5",
+            "stabilityai/stable-diffusion-2-1",
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            "CompVis/stable-diffusion-v1-4"
+        ]
+
+    def prepare_image_dataset(self, image_files, captions, resolution=512):
+        dataset = []
+        for img_path, caption in zip(image_files, captions):
+            try:
+                image = Image.open(img_path).convert("RGB")
+                dataset.append({"image": image, "caption": caption, "image_path": img_path})
+            except:
+                continue
+        return dataset
+
+    def simulate_training(self, job_id, model_name, dataset, r=16, lora_alpha=32, lora_dropout=0.1, num_epochs=10):
+        self.training_jobs[job_id]["status"] = "training"
+        self.training_jobs[job_id]["progress"] = 0
+        total_steps = num_epochs * len(dataset)
+        for step in range(total_steps):
+            time.sleep(0.2)
+            self.training_jobs[job_id]["progress"] = int((step + 1) / total_steps * 100)
+        # Simular criação do modelo
+        output_dir = f"./lora_models/{job_id}"
+        os.makedirs(output_dir, exist_ok=True)
+        with open(f"{output_dir}/adapter_model.safetensors", "w") as f:
+            f.write("Simulated LoRA model")
+        with open(f"{output_dir}/adapter_config.json", "w") as f:
+            json.dump({"model_name": model_name}, f)
+        self.training_jobs[job_id]["status"] = "completed"
+        self.training_jobs[job_id]["model_path"] = output_dir
+        self.training_jobs[job_id]["progress"] = 100
+
+    def start_training(self, model_name, image_files, captions, **kwargs):
+        job_id = str(uuid.uuid4())
+        dataset = self.prepare_image_dataset(image_files, captions)
+        self.training_jobs[job_id] = {"id": job_id, "status": "queued", "progress": 0, "model_name": model_name, "model_path": None}
+        thread = threading.Thread(target=self.simulate_training, args=(job_id, model_name, dataset), kwargs=kwargs)
+        thread.start()
+        return job_id
+
+    def get_training_status(self, job_id):
+        return self.training_jobs.get(job_id, {"error": "Job não encontrado"})
+
+    def list_trained_models(self):
+        models = []
+        lora_models_dir = Path("./lora_models")
+        if lora_models_dir.exists():
+            for model_dir in lora_models_dir.iterdir():
+                if model_dir.is_dir():
+                    models.append({"id": model_dir.name, "path": str(model_dir)})
+        return models
+
+    def create_download_zip(self, model_path):
+        zip_path = f"{model_path}.zip"
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for file_path in Path(model_path).rglob("*"):
+                if file_path.is_file():
+                    zipf.write(file_path, arcname=file_path.name)
+        return zip_path
+
+# --- Instância ---
+trainer = LoRAImageTrainer()
+
+# --- Gradio Interface ---
+def start_training_wrapper(model_name, files, captions_text, trigger_word, r, lora_alpha, lora_dropout, num_epochs, learning_rate, batch_size, resolution):
+    if not files or len(files) < 3:
+        return "❌ Envie pelo menos 3 imagens para treinamento!"
+    image_files = [f.name for f in files]
+    captions = [line.strip() for line in captions_text.split("\n") if line.strip()]
+    while len(captions) < len(files):
+        captions.append(trigger_word or f"image {len(captions)+1}")
+    captions = captions[:len(files)]
+    job_id = trainer.start_training(model_name, image_files, captions)
+    return f"✅ Treinamento iniciado! Job ID: {job_id}"
+
+def check_status_wrapper(job_id):
+    status = trainer.get_training_status(job_id.strip())
+    if "error" in status:
+        return status["error"]
+    return f"Status: {status['status']}\nProgresso: {status['progress']}%"
+
+def list_models_wrapper():
+    models = trainer.list_trained_models()
+    if not models:
+        return "📭 Nenhum modelo encontrado."
+    text = ""
+    for m in models:
+        text += f"ID: {m['id']}\nPath: {m['path']}\n---\n"
+    return text
+
+def download_model_wrapper(job_id):
+    status = trainer.get_training_status(job_id.strip())
+    if status.get("status") != "completed":
+        return None, "❌ Modelo não disponível ou treinamento não concluído."
+    zip_path = trainer.create_download_zip(status["model_path"])
+    return zip_path, "✅ Clique para baixar"
+
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎨 LoRA Image Trainer")
-    
-    with gr.Row():
-        with gr.Column():
-            images_input = gr.File(label="Selecione imagens", file_types=[".jpg", ".png"], file_types_allow_multiple=True)
-            model_name_input = gr.Textbox(label="Modelo Base", value="Stable Diffusion 1.5")
-            trigger_word_input = gr.Textbox(label="Trigger Word", placeholder="Ex: myStyle")
-        with gr.Column():
-            epochs_input = gr.Slider(5, 50, value=10, step=1, label="Épocas")
-            batch_size_input = gr.Slider(1, 8, value=1, step=1, label="Batch Size")
-            learning_rate_input = gr.Number(value=1e-4, label="Learning Rate")
-            lora_r_input = gr.Slider(4, 128, value=16, step=1, label="LoRA Rank (r)")
-            lora_alpha_input = gr.Slider(1, 128, value=32, step=1, label="LoRA Alpha")
-            lora_dropout_input = gr.Slider(0.0, 0.5, value=0.1, step=0.01, label="LoRA Dropout")
-    
-    train_button = gr.Button("Iniciar Treinamento")
-    output_text = gr.Textbox(label="Status do Treinamento")
-
-    train_button.click(
-        fn=start_training,
-        inputs=[
-            images_input, model_name_input, trigger_word_input,
-            epochs_input, batch_size_input, learning_rate_input,
-            lora_r_input, lora_alpha_input, lora_dropout_input
-        ],
-        outputs=output_text
-    )
-
-# NÃO chamar demo.launch() aqui
-# Hugging Face detecta automaticamente a variável `demo`
+    with gr.Tab("🎯 Treinar LoRA"):
+        model_dropdown = gr.Dropdown(choices=trainer.get_available_models(), value="runwayml/stable-diffusion-v1-5", label="Modelo Base")
+        image_files = gr.File(file_types=["image"], file_count="multiple", label="Imagens")
+        trigger_word = gr.Textbox(label="Trigger Word")
+        captions_text = gr.Textbox(label="Legendas (opcional)")
+        train_button = gr.Button("Iniciar Treinamento")
+        train_output = gr.Textbox(label="Resultado")
+        train_button.click(start_training_wrapper, inputs=[model_dropdown, image_files, captions_text, trigger_word, 16, 32, 0.1, 10, 0.0001, 1, 512], outputs=train_output)
+
+    with gr.Tab("📊 Status do Treinamento"):
+        job_id_input = gr.Textbox(label="Job ID")
+        status_button = gr.Button("Verificar Status")
+        status_output = gr.Textbox(label="Status")
+        status_button.click(check_status_wrapper, inputs=job_id_input, outputs=status_output)
+
+    with gr.Tab("📚 Modelos e Download"):
+        list_button = gr.Button("Listar Modelos")
+        models_output = gr.Textbox(label="Modelos")
+        list_button.click(list_models_wrapper, outputs=models_output)
+        download_job_id = gr.Textbox(label="Job ID para Download")
+        download_button = gr.Button("Preparar Download")
+        download_file = gr.File(label="Download do Modelo")
+        download_status = gr.Textbox(label="Status Download")
+        download_button.click(download_model_wrapper, inputs=download_job_id, outputs=[download_file, download_status])
+
+if __name__ == "__main__":
+    os.makedirs("./lora_models", exist_ok=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)