Spaces:
Running on Zero
Running on Zero
| import torch | |
| from torch.utils.data import DataLoader | |
| from datasets import load_dataset | |
| from torchvision import transforms | |
| from huggingface_hub import HfApi, create_repo, upload_folder | |
| import os | |
| from tqdm.auto import tqdm | |
| import json | |
| class LouTrainer: | |
| def __init__(self, model, diffusion, device='cuda'): | |
| self.model = model | |
| self.diffusion = diffusion | |
| self.device = device | |
| self.optimizer = None | |
| self.current_epoch = 0 | |
| def setup_optimizer(self, lr=1e-4): | |
| self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=0.01) | |
| self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=1000) | |
| def prepare_dataset(self, dataset_name="diffusers/pokemon-gpt4-captions", batch_size=4, image_size=256): | |
| """Charge et prépare le dataset""" | |
| # Transformations | |
| transform = transforms.Compose([ | |
| transforms.Resize(image_size), | |
| transforms.CenterCrop(image_size), | |
| transforms.RandomHorizontalFlip(), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) | |
| ]) | |
| # Chargement dataset | |
| dataset = load_dataset(dataset_name, split="train") | |
| def transform_examples(examples): | |
| images = [transform(img.convert("RGB")) for img in examples["image"]] | |
| return {"pixel_values": images} | |
| dataset.set_transform(transform_examples) | |
| dataloader = DataLoader( | |
| dataset, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| num_workers=0, | |
| pin_memory=True | |
| ) | |
| return dataloader | |
| def train_epoch(self, dataloader, progress_bar=None): | |
| """Entraîne une epoch""" | |
| total_loss = 0 | |
| for batch_idx, batch in enumerate(dataloader): | |
| images = batch["pixel_values"].to(self.device) | |
| self.optimizer.zero_grad() | |
| loss = self.diffusion.training_step(images) | |
| loss.backward() | |
| # Gradient clipping | |
| torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) | |
| self.optimizer.step() | |
| total_loss += loss.item() | |
| if progress_bar: | |
| progress_bar.update(1) | |
| progress_bar.set_postfix({"loss": f"{loss.item():.4f}"}) | |
| self.scheduler.step() | |
| return total_loss / len(dataloader) | |
| def train(self, dataset_name, num_epochs, lr, batch_size=4, image_size=256, | |
| save_every=5, repo_id=None, token=None, progress_callback=None): | |
| """Entraînement complet""" | |
| self.setup_optimizer(lr) | |
| dataloader = self.prepare_dataset(dataset_name, batch_size, image_size) | |
| total_steps = len(dataloader) * num_epochs | |
| with tqdm(total=total_steps, desc="Entraînement Lou") as pbar: | |
| for epoch in range(num_epochs): | |
| self.current_epoch = epoch | |
| avg_loss = self.train_epoch(dataloader, pbar) | |
| print(f"Epoch {epoch+1}/{num_epochs} - Loss: {avg_loss:.4f}") | |
| # Sauvegarde périodique | |
| if (epoch + 1) % save_every == 0 and repo_id and token: | |
| self.save_checkpoint(repo_id, token, epoch + 1) | |
| if progress_callback: | |
| progress_callback(epoch + 1, num_epochs, avg_loss) | |
| # Sauvegarde finale | |
| if repo_id and token: | |
| self.save_checkpoint(repo_id, token, "final") | |
| return True | |
| def save_checkpoint(self, repo_id, token, epoch): | |
| """Sauvegarde sur HF Hub""" | |
| os.makedirs("checkpoints", exist_ok=True) | |
| # Sauvegarde locale | |
| checkpoint_path = f"checkpoints/lou_epoch_{epoch}.pt" | |
| torch.save({ | |
| 'epoch': epoch, | |
| 'model_state_dict': self.model.state_dict(), | |
| 'optimizer_state_dict': self.optimizer.state_dict() if self.optimizer else None, | |
| }, checkpoint_path) | |
| # Upload vers HF | |
| try: | |
| api = HfApi(token=token) | |
| try: | |
| create_repo(repo_id, private=False, token=token) | |
| except: | |
| pass | |
| api.upload_file( | |
| path_or_fileobj=checkpoint_path, | |
| path_in_repo=f"lou_epoch_{epoch}.pt", | |
| repo_id=repo_id, | |
| token=token | |
| ) | |
| print(f"✓ Checkpoint epoch {epoch} uploadé sur {repo_id}") | |
| except Exception as e: | |
| print(f"✗ Erreur upload: {e}") | |
| def load_checkpoint(self, checkpoint_path): | |
| """Charge un checkpoint""" | |
| checkpoint = torch.load(checkpoint_path, map_location=self.device) | |
| self.model.load_state_dict(checkpoint['model_state_dict']) | |
| if self.optimizer and checkpoint['optimizer_state_dict']: | |
| self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) | |
| self.current_epoch = checkpoint.get('epoch', 0) | |
| print(f"✓ Checkpoint chargé (epoch {self.current_epoch})") |