Delete app.py

app.py

@@ -1,243 +0,0 @@
-
-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,
-    resolution,
-    learning_rate,
-    batch_size,
-    num_epochs,
-    train_prompt,
-):
-    if dataset_zip 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"
-    zip_path = os.path.join("./data", os.path.basename(dataset_zip.name))
-    with open(zip_path, "wb") as f:
-        f.write(dataset_zip.read())
-
-    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)
-
-