# ========================================== # 0. INSTALLATION DES DÉPENDANCES ULTRA-PRO # ========================================== !pip install datasets torch torchvision pillow huggingface_hub import os import shutil import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader from torchvision import transforms, models from datasets import load_dataset from PIL import Image from google.colab import files # ========================================== # 1. CONFIGURATION CHIRURGICALE # ========================================== DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") MODEL_INPUT_SIZE = 224 # Taille requise par le Vision Transformer (ViT) SKIN_SIZE = 128 # Taille cible HD du skin Minecraft MODEL_PATH = "vit_minecraft_core.pth" # ========================================== # 2. ARCHITECTURE : VISION TRANSFORMER # ========================================== class ViTSkinGenerator(nn.Module): def __init__(self): super().__init__() # Modèle de base Google pré-entraîné vit = models.vit_b_16(pretrained=True) self.transformer_core = vit num_features = vit.heads.head.in_features # Décodeur linéaire profond pour projeter l'attention sur la texture 128x128x3 self.transformer_core.heads = nn.Sequential( nn.Linear(num_features, 2048), nn.GELU(), nn.Dropout(0.1), nn.Linear(2048, 4096), nn.GELU(), nn.Linear(4096, 3 * SKIN_SIZE * SKIN_SIZE), nn.Sigmoid() ) def forward(self, x): out = self.transformer_core(x) out = out.view(-1, 3, SKIN_SIZE, SKIN_SIZE) return out class MinecraftTransformerDataset(Dataset): def __init__(self, img_dir, skin_dir): self.img_dir = img_dir self.skin_dir = skin_dir self.filenames = sorted(os.listdir(img_dir)) self.transform_img = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) self.transform_skin = transforms.Compose([transforms.ToTensor()]) def __len__(self): return len(self.filenames) def __getitem__(self, idx): f = self.filenames[idx] return self.transform_img(Image.open(os.path.join(self.img_dir, f)).convert("RGB")), \ self.transform_skin(Image.open(os.path.join(self.skin_dir, f)).convert("RGB")) # ========================================== # 3. PIPELINE D'ENTRAÎNEMENT (TRAIN) # ========================================== def train_pipeline(epochs=35, batch_size=16, lr=3e-4, num_samples=437): """Télécharge le dataset, initialise le modèle (ou charge l'existant) et lance l'entraînement.""" print(f"🚀 Initialisation de l'entraînement sur : {DEVICE}") # Nettoyage du cache pour éviter les erreurs de taille d'image if os.path.exists("dataset"): shutil.rmtree("dataset") os.makedirs("dataset/images", exist_ok=True) os.makedirs("dataset/skins", exist_ok=True) print("📥 Extraction du dataset Hugging Face...") hf_dataset = load_dataset("zhoudoe23/minecraft-skins-1.1m-prerendered", split="train", streaming=True) count = 0 for item in hf_dataset: if count >= num_samples: break try: render_img = item["rendering_3d"].convert("RGB").resize((MODEL_INPUT_SIZE, MODEL_INPUT_SIZE)) skin_img = item["raw_skin"].convert("RGB").resize((SKIN_SIZE, SKIN_SIZE), Image.NEAREST) render_img.save(f"dataset/images/skin_{count}.png") skin_img.save(f"dataset/skins/skin_{count}.png") count += 1 except: continue dataset = MinecraftTransformerDataset("dataset/images", "dataset/skins") dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) model = ViTSkinGenerator().to(DEVICE) optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) criterion = nn.SmoothL1Loss() if os.path.exists(MODEL_PATH): print(f"💾 Modèle existant trouvé. Reprise de l'apprentissage depuis {MODEL_PATH}...") model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE)) print("\n🏋️ Alignement du Vision Transformer...") for epoch in range(epochs): model.train() epoch_loss = 0 for imgs, skins in dataloader: imgs, skins = imgs.to(DEVICE), skins.to(DEVICE) outputs = model(imgs) loss = criterion(outputs, skins) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss += loss.item() print(f"Époque [{epoch+1}/{epochs}] | Indice de précision (Loss): {epoch_loss/len(dataloader):.6f}") torch.save(model.state_dict(), MODEL_PATH) print(f"✅ Modèle sauvegardé avec succès dans : {MODEL_PATH}") # ========================================== # 4. PIPELINE D'INFÉRENCE (USE) # ========================================== def use_pipeline(image_path, output_path="skin_vit_ultimate.png", download=True): """Prend n'importe quelle image d'entrée, charge les poids du modèle et génère le skin.""" if not os.path.exists(MODEL_PATH): print(f"❌ Erreur : Le fichier de poids '{MODEL_PATH}' n'existe pas. Lancez d'abord l'entraînement.") return # Initialisation et chargement des poids de manière déterministe model = ViTSkinGenerator().to(DEVICE) model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE)) model.eval() if not os.path.exists(image_path): print(f"⚠️ Image d'entrée '{image_path}' introuvable.") return # Traitement de l'image d'entrée img = Image.open(image_path).convert("RGB").resize((MODEL_INPUT_SIZE, MODEL_INPUT_SIZE), Image.Resampling.LANCZOS) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) tensor_in = transform(img).unsqueeze(0).to(DEVICE) with torch.no_grad(): tensor_out = model(tensor_in).squeeze(0).cpu() # Transformation finale en Pixel Art strict (128x128) tensor_out = torch.clamp(tensor_out, 0, 1) skin_pil = transforms.ToPILImage()(tensor_out) skin_final = skin_pil.resize((SKIN_SIZE, SKIN_SIZE), Image.NEAREST) skin_final.save(output_path) print(f"✨ Skin extrait avec succès dans : {output_path}") if download: files.download(output_path) # ========================================== # 5. EXEMPLES D'UTILISATION (À MODIFIER SELON VOS ENVIES) # ========================================== # Étape A : Pour lancer l'entraînement (Décommentez la ligne ci-dessous pour l'activer) train_pipeline(epochs=25, batch_size=16, num_samples=437) # Étape B : Pour générer un skin à partir d'une image personnalisée (Décommentez ci-dessous) # use_pipeline(image_path="A.jpg", output_path="resultat_final.png")