Update app.py

app.py

@@ -1,24 +1,67 @@
-
 import gradio as gr
 import os
 import torch
 from accelerate import Accelerator
-from accelerate.utils import set_seed
-from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler, StableDiffusionPipeline
+from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
 from diffusers.optimization import get_scheduler
-from diffusers.training_utils import EMAModel
-from diffusers.models.attention_processor import LoRAAttnProcessor as DiffusersLoRAAttnProcessor
-from huggingface_hub import create_repo, upload_folder
 from PIL import Image
 from torch.utils.data import Dataset
 from torchvision import transforms
-from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer
 import zipfile
 import shutil
 from safetensors.torch import save_file
-
-# Placeholder para o script de treinamento
+import torch.nn as nn
+
+# Função para criar camadas LoRA
+def create_lora_layers(module, rank=4):
+    if isinstance(module, nn.Linear):
+        lora_down = nn.Linear(module.in_features, rank, bias=False)
+        lora_up = nn.Linear(rank, module.out_features, bias=False)
+        nn.init.zeros_(lora_up.weight)  # Inicialização zero para começar neutro
+        return lora_down, lora_up
+    return None, None
+
+# Dataset simplificado
+class DreamBoothDataset(Dataset):
+    def __init__(self, instance_data_root, tokenizer, size=512, train_prompt="a photo of sks dog"):
+        self.instance_data_root = instance_data_root
+        self.tokenizer = tokenizer
+        self.size = size
+        self.train_prompt = train_prompt
+        self.instance_images_path = [
+            os.path.join(instance_data_root, file_path)
+            for file_path in os.listdir(instance_data_root)
+            if file_path.endswith((".png", ".jpg", ".jpeg"))
+        ]
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+    def __len__(self):
+        return len(self.instance_images_path)
+
+    def __getitem__(self, index):
+        instance_image = Image.open(self.instance_images_path[index])
+        if not instance_image.mode == "RGB":
+            instance_image = instance_image.convert("RGB")
+        example = {}
+        example["instance_images"] = self.transform(instance_image)
+        example["instance_prompt_ids"] = self.tokenizer(
+            self.train_prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids[0]
+        return example
+
+# Função principal de treinamento
 def train_lora(
     instance_data_dir: str,
     output_dir: str,
@@ -35,40 +78,57 @@ def train_lora(
         mixed_precision="fp16",
     )
 
-    # Carregar tokenizer e modelo base
-    tokenizer = CLIPTokenizer.from_pretrained(
-        pretrained_model_name_or_path, subfolder="tokenizer"
-    )
-    text_encoder = CLIPTextModel.from_pretrained(
-        pretrained_model_name_or_path, subfolder="text_encoder"
-    )
-    vae = AutoencoderKL.from_pretrained(
-        pretrained_model_name_or_path, subfolder="vae"
-    )
-    unet = UNet2DConditionModel.from_pretrained(
-        pretrained_model_name_or_path, subfolder="unet"
-    )
+    # Carregar modelos
+    tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder")
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
 
-    # Congelar parâmetros do VAE e Text Encoder
+    # Congelar VAE e Text Encoder
     vae.requires_grad_(False)
     text_encoder.requires_grad_(False)
-
-    # Configurar LoRA
-    # Adicionar adaptadores LoRA ao UNet
-    # A função `add_adapter` do diffusers já configura os módulos LoRA e os torna treináveis.
-    unet.add_adapter(DiffusersLoRAAttnProcessor)
+    unet.requires_grad_(False)
+
+    # Injetar LoRA no UNet
+    lora_layers = []
+    for name, module in unet.named_modules():
+        if name.endswith("to_q") or name.endswith("to_k") or name.endswith("to_v") or name.endswith("to_out.0"):
+            lora_down, lora_up = create_lora_layers(module, rank=4)
+            if lora_down is not None:
+                module.lora_down = lora_down.to(module.weight.device)
+                module.lora_up = lora_up.to(module.weight.device)
+                lora_layers.extend([module.lora_down, module.lora_up])
+
+                # Guardar forward original
+                if not hasattr(module, "_original_forward"):
+                    module._original_forward = module.forward
+
+                # Criar novo forward com LoRA
+                def forward_with_lora(self, x):
+                    original_output = self._original_forward(x)
+                    lora_output = self.lora_up(self.lora_down(x))
+                    return original_output + lora_output
+
+                # Associar o novo forward ao módulo
+                import types
+                module.forward = types.MethodType(forward_with_lora, module)
+
+    # Liberar apenas parâmetros LoRA
+    for layer in lora_layers:
+        layer.requires_grad_(True)
+
+    # Coletar parâmetros treináveis
+    lora_parameters = []
+    for layer in lora_layers:
+        lora_parameters.extend(layer.parameters())
 
     # Otimizador
-    # Apenas os parâmetros do LoRA devem ser treináveis
-    # O `add_adapter` já faz isso, então podemos simplesmente pegar os parâmetros treináveis do UNet.
-    lora_parameters = list(filter(lambda p: p.requires_grad, unet.parameters()))
+    optimizer = torch.optim.AdamW(lora_parameters, lr=learning_rate)
 
-    optimizer = torch.optim.AdamW(
-        lora_parameters,
-        lr=learning_rate,
-    )
+    # Scheduler de ruído
+    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
 
-    # Scheduler
+    # Scheduler de learning rate
     lr_scheduler = get_scheduler(
         "constant",
         optimizer=optimizer,
@@ -76,92 +136,61 @@ def train_lora(
         num_training_steps=num_epochs * len(os.listdir(instance_data_dir)),
     )
 
-    # Dataset e DataLoader (simplificado para o exemplo)
-    class DreamBoothDataset(Dataset):
-        def __init__(self, instance_data_root, tokenizer, size=512, train_prompt="a photo of sks dog"):
-            self.instance_data_root = instance_data_root
-            self.tokenizer = tokenizer
-            self.size = size
-            self.train_prompt = train_prompt
-            self.instance_images_path = [os.path.join(instance_data_root, file_path) for file_path in os.listdir(instance_data_root) if file_path.endswith((".png", ".jpg", ".jpeg"))]
-            self.transform = transforms.Compose(
-                [
-                    transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
-                    transforms.CenterCrop(size),
-                    transforms.ToTensor(),
-                    transforms.Normalize([0.5], [0.5]),
-                ]
-            )
-
-        def __len__(self):
-            return len(self.instance_images_path)
-
-        def __getitem__(self, index):
-            instance_image = Image.open(self.instance_images_path[index])
-            if not instance_image.mode == "RGB":
-                instance_image = instance_image.convert("RGB")
-            example = {}
-            example["instance_images"] = self.transform(instance_image)
-            example["instance_prompt_ids"] = self.tokenizer(self.train_prompt,
-                                                            truncation=True,
-                                                            padding="max_length",
-                                                            max_length=self.tokenizer.model_max_length,
-                                                            return_tensors="pt",
-                                                            ).input_ids[0]
-            return example
-
+    # Dataset e DataLoader
     train_dataset = DreamBoothDataset(instance_data_dir, tokenizer, resolution, train_prompt)
     train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
 
-    # Preparar para treinamento com Accelerator
+    # Preparar com Accelerator
     unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
         unet, optimizer, train_dataloader, lr_scheduler
     )
 
-    # Loop de treinamento
+    # Treinamento
+    global_step = 0
     for epoch in range(num_epochs):
         unet.train()
         for step, batch in enumerate(train_dataloader):
             with accelerator.accumulate(unet):
-                # Forward pass
-                latents = vae.encode(batch["instance_images"]).latent_dist.sample()
+                # Preparar dados
+                pixel_values = batch["instance_images"].to(accelerator.device)
+                latents = vae.encode(pixel_values).latent_dist.sample()
                 latents = latents * vae.config.scaling_factor
 
-                noise = torch.randn_like(latents)
-                timesteps = torch.randint(0, 1000, (batch_size,), device=latents.device).long()
+                noise = torch.randn_like(latents).to(accelerator.device)
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=latents.device).long()
 
-                noisy_latents = DDPMScheduler().add_noise(latents, noise, timesteps)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
-                encoder_hidden_states = text_encoder(batch["instance_prompt_ids"])[0]
+                encoder_hidden_states = text_encoder(batch["instance_prompt_ids"].to(accelerator.device))[0]
 
+                # Predição
                 model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
-                # Calcular perda
+                # Perda
                 loss = torch.nn.functional.mse_loss(model_pred.float(), noise.float(), reduction="mean")
 
-                # Backward pass
+                # Backprop
                 accelerator.backward(loss)
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
 
-            accelerator.log({"loss": loss.item()}, step=epoch * len(train_dataloader) + step)
-            print(f"Epoch {epoch}, Step {step}, Loss: {loss.item()}")
+            global_step += 1
+            print(f"Epoch {epoch + 1}/{num_epochs}, Step {step + 1}, Loss: {loss.item():.6f}")
 
-    # Salvar o modelo treinado
-    # Salvar apenas os pesos LoRA
-    lora_state_dict = {} 
-    for name, param in unet.named_parameters():
-        if "lora" in name:
-            lora_state_dict[name] = param
+    # Salvar LoRA
+    lora_state_dict = {}
+    for name, module in unet.named_modules():
+        if hasattr(module, "lora_down") and hasattr(module, "lora_up"):
+            lora_state_dict[f"{name}.lora_down.weight"] = module.lora_down.weight
+            lora_state_dict[f"{name}.lora_up.weight"] = module.lora_up.weight
 
     lora_path = os.path.join(output_dir, "lora_model.safetensors")
-    # Usar safetensors para salvar o modelo
     save_file(lora_state_dict, lora_path)
 
     return lora_path
 
-
+# Função para Gradio
 def run_training(
     dataset_zip_file,
     resolution,
@@ -181,62 +210,51 @@ def run_training(
     os.makedirs("./data/dataset", exist_ok=True)
     os.makedirs("./outputs", exist_ok=True)
 
-    # Salvar e extrair o dataset
+    # Extrair dataset
     dataset_dir = "./data/dataset"
-    # O objeto dataset_zip_file do Gradio tem um atributo .name que é o caminho do arquivo temporário
     zip_path = dataset_zip_file.name
-    
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
         zip_ref.extractall(dataset_dir)
 
-    # Iniciar treinamento
+    # Treinar
     output_dir = "./outputs"
-    
-    lora_model_path = train_lora(
-        instance_data_dir=dataset_dir,
-        output_dir=output_dir,
-        resolution=resolution,
-        learning_rate=learning_rate,
-        batch_size=batch_size,
-        num_epochs=num_epochs,
-        train_prompt=train_prompt,
-    )
-
-    return f"Treinamento concluído! Modelo salvo em: {lora_model_path}", lora_model_path
-
-
+    try:
+        lora_model_path = train_lora(
+            instance_data_dir=dataset_dir,
+            output_dir=output_dir,
+            resolution=resolution,
+            learning_rate=learning_rate,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            train_prompt=train_prompt,
+        )
+        return f"✅ Treinamento concluído! Modelo salvo em: {lora_model_path}", lora_model_path
+    except Exception as e:
+        return f"❌ Erro durante o treinamento: {str(e)}", None
+
+# Interface Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Treinador LoRA para Hugging Face Spaces")
+    gr.Markdown("# 🧠 Treinador LoRA para Stable Diffusion")
 
     with gr.Row():
         with gr.Column():
-            dataset_zip = gr.File(label="Upload do Dataset (ZIP)", file_types=[".zip"])
-            resolution = gr.Slider(minimum=128, maximum=1024, value=512, step=128, label="Resolução da Imagem")
-            learning_rate = gr.Number(value=1e-4, label="Learning Rate")
-            batch_size = gr.Slider(minimum=1, maximum=8, value=1, step=1, label="Batch Size")
-            num_epochs = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Número de Epochs")
-            train_prompt = gr.Textbox(label="Prompt de Treinamento (ex: a photo of sks dog)", value="a photo of sks dog")
-            train_button = gr.Button("Iniciar Treinamento")
+            dataset_zip = gr.File(label="📁 Upload do Dataset (ZIP)", file_types=[".zip"])
+            resolution = gr.Slider(minimum=128, maximum=1024, value=512, step=128, label="📏 Resolução da Imagem")
+            learning_rate = gr.Number(value=1e-4, label="📈 Learning Rate")
+            batch_size = gr.Slider(minimum=1, maximum=8, value=1, step=1, label="📦 Batch Size")
+            num_epochs = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="🔁 Número de Epochs")
+            train_prompt = gr.Textbox(label="📝 Prompt de Treinamento (ex: a photo of sks dog)", value="a photo of sks dog")
+            train_button = gr.Button("🚀 Iniciar Treinamento", variant="primary")
 
         with gr.Column():
-            output_text = gr.Textbox(label="Status do Treinamento")
-            output_file = gr.File(label="Modelo LoRA Treinado")
+            output_text = gr.Textbox(label="📊 Status do Treinamento", lines=5)
+            output_file = gr.File(label="💾 Modelo LoRA Treinado")
 
     train_button.click(
         run_training,
-        inputs=[
-            dataset_zip,
-            resolution,
-            learning_rate,
-            batch_size,
-            num_epochs,
-            train_prompt,
-        ],
+        inputs=[dataset_zip, resolution, learning_rate, batch_size, num_epochs, train_prompt],
         outputs=[output_text, output_file],
     )
 
-
 if __name__ == "__main__":
-    demo.launch(debug=True)
-
-
+    demo.launch(debug=True)