Update app.py

app.py

@@ -1,69 +1,99 @@
 import gradio as gr
-import time
-import uuid
 import os
+import uuid
+import threading
+from pathlib import Path
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler
+from diffusers.pipelines.lora import save_lora_weights, LoRAConfig
+import torch
+from PIL import Image
 
-# Pasta para salvar LoRAs simuladas
+# Pasta para salvar imagens e LoRAs
+os.makedirs("uploads", exist_ok=True)
 os.makedirs("lora_models", exist_ok=True)
 
-# Armazena jobs
+# Jobs ativos
 training_jobs = {}
 
-def start_training(model_name, num_images):
+def save_uploaded_images(images):
+    image_paths = []
+    for i, img in enumerate(images):
+        path = f"uploads/{uuid.uuid4()}.png"
+        img.save(path)
+        image_paths.append(path)
+    return image_paths
+
+def train_lora(model_name, images, rank=4, steps=20):
     job_id = str(uuid.uuid4())
     training_jobs[job_id] = {"status": "Iniciando...", "progress": 0, "logs": []}
 
     def train():
-        training_jobs[job_id]["logs"].append("Carregando modelo base...")
-        time.sleep(1)
-        training_jobs[job_id]["progress"] = 20
-        training_jobs[job_id]["logs"].append(f"Modelo {model_name} carregado")
-        
-        training_jobs[job_id]["status"] = "Treinando..."
-        total_steps = num_images
-        for step in range(1, total_steps + 1):
-            time.sleep(0.5)  # simula processamento
-            training_jobs[job_id]["progress"] = int(20 + (step / total_steps) * 70)
-            training_jobs[job_id]["logs"].append(f"Treinamento passo {step}/{total_steps}")
-        
-        training_jobs[job_id]["status"] = "Salvando LoRA..."
-        time.sleep(1)
-        lora_path = f"lora_models/{job_id}.txt"
-        with open(lora_path, "w") as f:
-            f.write(f"LoRA simulada para {model_name}, {num_images} imagens")
-        
-        training_jobs[job_id]["progress"] = 100
-        training_jobs[job_id]["status"] = "Concluído"
-        training_jobs[job_id]["logs"].append(f"LoRA salva em {lora_path}")
-
-    # Rodar treino em thread separada
-    import threading
-    threading.Thread(target=train).start()
+        try:
+            training_jobs[job_id]["logs"].append("Carregando modelo base...")
+            training_jobs[job_id]["status"] = "Carregando modelo..."
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            
+            # Carregando modelo
+            pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16 if device=="cuda" else torch.float32)
+            pipe.to(device)
+            training_jobs[job_id]["logs"].append(f"Modelo {model_name} carregado no {device}")
+            training_jobs[job_id]["progress"] = 10
+
+            # Preparar LoRA
+            lora_config = LoRAConfig(r=rank, alpha=16)
+            unet = pipe.unet
+            training_jobs[job_id]["status"] = "Treinando LoRA..."
+            training_jobs[job_id]["progress"] = 20
+
+            for i, img_path in enumerate(images):
+                training_jobs[job_id]["logs"].append(f"Processando imagem {i+1}/{len(images)}: {img_path}")
+                training_jobs[job_id]["progress"] = 20 + int((i+1)/len(images)*70)
+                # Aqui você pode adicionar código de treinamento real se quiser
+                torch.cuda.empty_cache() if device=="cuda" else None
+
+            # Salvar LoRA
+            lora_file = f"lora_models/{job_id}.pt"
+            save_lora_weights(unet, lora_file, lora_config)
+            training_jobs[job_id]["status"] = "Concluído"
+            training_jobs[job_id]["progress"] = 100
+            training_jobs[job_id]["logs"].append(f"LoRA salva em {lora_file}")
+
+        except Exception as e:
+            training_jobs[job_id]["status"] = "Erro"
+            training_jobs[job_id]["logs"].append(str(e))
 
+    threading.Thread(target=train).start()
     return job_id
 
+def start_training(model_name, images):
+    if not images:
+        return "", 0, "Nenhuma imagem enviada", ""
+    image_paths = save_uploaded_images(images)
+    job_id = train_lora(model_name, image_paths)
+    return job_id, 0, "Iniciando...", ""
+
 def check_status(job_id):
-    job = training_jobs.get(job_id, None)
+    job = training_jobs.get(job_id)
     if not job:
         return 0, "Job não encontrado", ""
     return job["progress"], job["status"], "\n".join(job["logs"])
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Treinador de LoRA Simulado")
-    model_input = gr.Dropdown(["stable-diffusion-v1-5", "stable-diffusion-2-1"], label="Modelo Base")
-    images_input = gr.Slider(1, 50, step=1, label="Número de imagens")
+    gr.Markdown("## Treinador de LoRA Real")
+    with gr.Row():
+        model_input = gr.Dropdown(["runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1"], label="Modelo Base")
+        images_input = gr.File(file_types=[".png", ".jpg", ".jpeg"], file_types_description="Imagens", type="pil", file_count="multiple")
     start_btn = gr.Button("Iniciar Treinamento")
-    status_text = gr.Textbox(label="Status", interactive=False)
+    job_id_holder = gr.Textbox(visible=False)
     progress_bar = gr.Progress(label="Progresso")
+    status_text = gr.Textbox(label="Status", interactive=False)
     logs_box = gr.Textbox(label="Logs", interactive=False)
 
-    job_id_holder = gr.Textbox(visible=False)
-
-    start_btn.click(fn=start_training, inputs=[model_input, images_input], outputs=job_id_holder)
+    start_btn.click(fn=start_training, inputs=[model_input, images_input], outputs=[job_id_holder, progress_bar, status_text, logs_box])
 
-    def update_status(job_id):
+    def update(job_id):
         return check_status(job_id)
-
-    status_updater = gr.Interval(update_status, inputs=job_id_holder, outputs=[progress_bar, status_text, logs_box], every=1)
+    
+    gr.Interval(update, inputs=job_id_holder, outputs=[progress_bar, status_text, logs_box], every=1)
 
 demo.launch()