Upload 2 files

app_1.py

@@ -0,0 +1,242 @@
+
+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.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
+def train_lora(
+    instance_data_dir: str,
+    output_dir: str,
+    resolution: int = 512,
+    learning_rate: float = 1e-4,
+    batch_size: int = 1,
+    num_epochs: int = 1,
+    train_prompt: str = "a photo of sks dog",
+    pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5",
+):
+    # Configurações básicas
+    accelerator = Accelerator(
+        gradient_accumulation_steps=1,
+        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"
+    )
+
+    # Congelar parâmetros do 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)
+
+    # 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,
+    )
+
+    # Scheduler
+    lr_scheduler = get_scheduler(
+        "constant",
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        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
+
+    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
+    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        unet, optimizer, train_dataloader, lr_scheduler
+    )
+
+    # Loop de treinamento
+    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()
+                latents = latents * vae.config.scaling_factor
+
+                noise = torch.randn_like(latents)
+                timesteps = torch.randint(0, 1000, (batch_size,), device=latents.device).long()
+
+                noisy_latents = DDPMScheduler().add_noise(latents, noise, timesteps)
+
+                encoder_hidden_states = text_encoder(batch["instance_prompt_ids"])[0]
+
+                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                # Calcular perda
+                loss = torch.nn.functional.mse_loss(model_pred.float(), noise.float(), reduction="mean")
+
+                # Backward pass
+                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()}")
+
+    # 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
+
+    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
+
+
+def run_training(
+    dataset_zip_file,
+    resolution,
+    learning_rate,
+    batch_size,
+    num_epochs,
+    train_prompt,
+):
+    if dataset_zip_file is None:
+        return "Por favor, faça o upload de um arquivo ZIP com seu dataset.", None
+
+    # Limpar diretórios anteriores
+    if os.path.exists("./data/dataset"):
+        shutil.rmtree("./data/dataset")
+    if os.path.exists("./outputs"):
+        shutil.rmtree("./outputs")
+    os.makedirs("./data/dataset", exist_ok=True)
+    os.makedirs("./outputs", exist_ok=True)
+
+    # Salvar e extrair o 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
+    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
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Treinador LoRA para Hugging Face Spaces")
+
+    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")
+
+        with gr.Column():
+            output_text = gr.Textbox(label="Status do Treinamento")
+            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,
+        ],
+        outputs=[output_text, output_file],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(debug=True)
+
+

requirements.txt

@@ -0,0 +1,13 @@
+torch
+accelerate
+transformers
+diffusers
+safetensors
+xformers
+gradio
+Pillow
+datasets
+
+
+torchvision
+