Update app.py

app.py

@@ -1,246 +1,69 @@
-import os
-import json
-import uuid
-import shutil
-import threading
+import gradio as gr
 import time
-from datetime import datetime
-from pathlib import Path
-from typing import Dict, List, Optional, Any, Tuple
-import zipfile
-import tempfile
+import uuid
+import os
 
-import gradio as gr
-import torch
-from PIL import Image
-import numpy as np
-from diffusers import (
-    StableDiffusionPipeline,
-    UNet2DConditionModel,
-    DDPMScheduler,
-    AutoencoderKL
-)
-from transformers import CLIPTextModel, CLIPTokenizer
-from peft import LoraConfig, get_peft_model, TaskType
-import logging
+# Pasta para salvar LoRAs simuladas
+os.makedirs("lora_models", exist_ok=True)
 
-# Configurar logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Armazena jobs
+training_jobs = {}
 
-class LoRAImageTrainer:
-    """Classe principal para treinamento de modelos LoRA para geração de imagens otimizada para baixo uso de GPU."""
-    
-    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) -> List[str]:
-        """Retorna lista de modelos base disponíveis para treinamento LoRA."""
-        return [
-            "runwayml/stable-diffusion-v1-5",
-            "stabilityai/stable-diffusion-2-1",
-            "stabilityai/stable-diffusion-xl-base-1.0",
-            "CompVis/stable-diffusion-v1-4"
-        ]
-    
-    def load_base_model(self, model_name: str):
-        """Carrega modelo base de difusão com otimizações para baixo uso de GPU."""
-        try:
-            if model_name in self.models_cache:
-                return self.models_cache[model_name]
-            
-            logger.info(f"Carregando modelo base: {model_name}")
-            
-            # Configurações para otimização de memória
-            model_kwargs = {
-                "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
-                "use_safetensors": True,
-                "variant": "fp16" if torch.cuda.is_available() else None,
-            }
-            
-            # Carregar pipeline completo
-            pipeline = StableDiffusionPipeline.from_pretrained(
-                model_name,
-                **model_kwargs
-            )
-            
-            if torch.cuda.is_available():
-                pipeline = pipeline.to(self.device)
-                # Habilitar attention slicing para economia de memória
-                pipeline.enable_attention_slicing()
-                # Habilitar memory efficient attention se disponível
-                try:
-                    pipeline.enable_xformers_memory_efficient_attention()
-                except:
-                    logger.warning("xformers não disponível, usando attention padrão")
-            
-            # Cache do modelo
-            self.models_cache[model_name] = pipeline
-            
-            return pipeline
-            
-        except Exception as e:
-            logger.error(f"Erro ao carregar modelo {model_name}: {str(e)}")
-            raise e
-    
-    def create_lora_config(self, 
-                          r: int = 16, 
-                          lora_alpha: int = 32, 
-                          lora_dropout: float = 0.1,
-                          target_modules: Optional[List[str]] = None) -> LoraConfig:
-        """Cria configuração LoRA otimizada para modelos de difusão."""
-        
-        if target_modules is None:
-            # Módulos padrão para UNet do Stable Diffusion
-            target_modules = [
-                "to_k", "to_q", "to_v", "to_out.0",
-                "proj_in", "proj_out",
-                "ff.net.0.proj", "ff.net.2"
-            ]
-        
-        return LoraConfig(
-            r=r,
-            lora_alpha=lora_alpha,
-            target_modules=target_modules,
-            lora_dropout=lora_dropout,
-            bias="none",
-            task_type=TaskType.DIFFUSION,
-        )
-    
-    def prepare_image_dataset(self, image_files: List[str], captions: List[str], resolution: int = 512) -> List[Dict]:
-        """Prepara dataset de imagens para treinamento."""
-        dataset = []
-        
-        for img_path, caption in zip(image_files, captions):
-            try:
-                # Carregar e redimensionar imagem
-                image = Image.open(img_path).convert("RGB")
-                
-                # Redimensionar mantendo aspect ratio
-                image = self.resize_image(image, resolution)
-                
-                dataset.append({
-                    "image": image,
-                    "caption": caption,
-                    "image_path": img_path
-                })
-                
-            except Exception as e:
-                logger.error(f"Erro ao processar imagem {img_path}: {str(e)}")
-                continue
-        
-        return dataset
-    
-    def resize_image(self, image: Image.Image, target_size: int) -> Image.Image:
-        """Redimensiona imagem mantendo aspect ratio e fazendo crop central se necessário."""
-        width, height = image.size
-        
-        # Calcular novo tamanho mantendo aspect ratio
-        if width > height:
-            new_width = target_size
-            new_height = int((height * target_size) / width)
-        else:
-            new_height = target_size
-            new_width = int((width * target_size) / height)
-        
-        # Redimensionar
-        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+def start_training(model_name, num_images):
+    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")
         
-        # Crop central para obter tamanho exato
-        if new_width != target_size or new_height != target_size:
-            left = (new_width - target_size) // 2
-            top = (new_height - target_size) // 2
-            right = left + target_size
-            bottom = top + target_size
-            
-            image = image.crop((left, top, right, bottom))
+        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}")
         
-        return image
-    
-    def simulate_training(self, 
-                         job_id: str,
-                         model_name: str,
-                         dataset: List[Dict],
-                         r: int = 16,
-                         lora_alpha: int = 32,
-                         lora_dropout: float = 0.1,
-                         num_epochs: int = 10,
-                         learning_rate: float = 1e-4,
-                         batch_size: int = 1,
-                         resolution: int = 512) -> None:
-        """Simula o processo de treinamento LoRA para imagens (versão demonstrativa)."""
+        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")
         
-        try:
-            # Atualizar status
-            self.training_jobs[job_id]["status"] = "loading_model"
-            self.training_jobs[job_id]["progress"] = 5
-            
-            # Simular carregamento do modelo base
-            time.sleep(2)
-            self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Modelo {model_name} carregado")
-            
-            # Preparar configuração LoRA
-            self.training_jobs[job_id]["status"] = "preparing_lora"
-            self.training_jobs[job_id]["progress"] = 15
-            time.sleep(1)
-            
-            lora_config = self.create_lora_config(r, lora_alpha, lora_dropout)
-            self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Configuração LoRA criada (r={r}, alpha={lora_alpha})")
-            
-            # Preparar dataset
-            self.training_jobs[job_id]["status"] = "preparing_data"
-            self.training_jobs[job_id]["progress"] = 25
-            time.sleep(1)
-            
-            self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Dataset preparado com {len(dataset)} imagens")
-            
-            # Simular treinamento
-            self.training_jobs[job_id]["status"] = "training"
-            self.training_jobs[job_id]["progress"] = 30
-            
-            total_steps = num_epochs * len(dataset)
-            current_step = 0
-            
-            for epoch in range(num_epochs):
-                for batch_idx in range(len(dataset)):
-                    current_step += 1
-                    
-                    # Simular tempo de processamento
-                    time.sleep(0.5)
-                    
-                    # Atualizar progresso
-                    progress = 30 + int((current_step / total_steps) * 60)
-                    self.training_jobs[job_id]["progress"] = min(progress, 90)
-                    
-                    # Simular loss decrescente
-                    loss = 0.8 - (current_step / total_steps) * 0.6
-                    
-                    if current_step % 5 == 0:  # Log a cada 5 steps
-                        log_message = f"Época {epoch+1}/{num_epochs}, Step {current_step}/{total_steps} - Loss: {loss:.4f}"
-                        self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - {log_message}")
-            
-            # Salvar modelo LoRA
-            self.training_jobs[job_id]["status"] = "saving"
-            self.training_jobs[job_id]["progress"] = 95
-            time.sleep(1)
-            
-            output_dir = f"./lora_models/{job_id}"
-            os.makedirs(output_dir, exist_ok=True)
-            
-            # Criar arquivos simulados do LoRA
-            lora_config_dict = {
-                "r": r,
-                "lora_alpha": lora_alpha,
-                "target_modules": ["to_k", "to_q", "to_v", "to_out.0"],
-                "lora_dropout": lora_dropout,
-                "bias": "none",
-                "task_type": "DIFFUSION",
-                "base_model_name": model_name,
-                "training_info": {
-                    "num_epochs": num_epochs,
-                    "learning_rate": learning_rate,
-                    "batch_size": batch_size,
-                    "resolution": resolution,
-                    "num_images": len(dataset)
+        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()
+
+    return job_id
+
+def check_status(job_id):
+    job = training_jobs.get(job_id, None)
+    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")
+    start_btn = gr.Button("Iniciar Treinamento")
+    status_text = gr.Textbox(label="Status", interactive=False)
+    progress_bar = gr.Progress(label="Progresso")
+    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)
+
+    def update_status(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)
+
+demo.launch()