Create main.py

main.py

@@ -0,0 +1,171 @@
+import torch
+import os
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from diffusers import UNet2DConditionModel, AutoencoderKL, DDPMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+from huggingface_hub import HfApi
+from torch.optim import AdamW
+from tqdm import tqdm
+import gc
+from torch.cuda.amp import autocast
+
+# Setare configurare CUDA pentru a reduce fragmentarea memoriei
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+# Verifică dacă GPU-ul este detectat
+print(torch.cuda.is_available())  
+
+img_dir = '/media/andrei_ursu/storage2/chess/branches/chessgpt/backend/src/experiments/full/primulTest/SD21data'
+
+# Definirea dataset-ului
+class ManualCaptionDataset(Dataset):
+    def __init__(self, img_dir, transform=None):
+        self.img_dir = img_dir
+        self.img_names = os.listdir(img_dir)
+        self.transform = transform
+        self.captions = []
+        
+        # Introducem manual descrierile pentru fiecare imagine
+        for img_name in self.img_names:
+            caption = 'Photo of Andrei smiling and dressed in winter clothes at a Christmas market'
+            self.captions.append(caption)
+        
+    def __len__(self):
+        return len(self.img_names)
+    
+    def __getitem__(self, idx):
+        img_name = os.path.join(self.img_dir, self.img_names[idx])
+        image = Image.open(img_name).convert("RGB")
+        caption = self.captions[idx]
+        
+        if self.transform:
+            image = self.transform(image)
+        
+        return image, caption
+
+# Configurare transformări
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),  # Dimensiune imagine redusă
+    transforms.ToTensor(),
+])
+
+# Crearea dataset-ului
+dataset = ManualCaptionDataset(img_dir=img_dir, transform=transform)
+dataloader = DataLoader(dataset, batch_size=1, shuffle=True)  # Dimensiune batch redusă
+
+# Încărcare model UNet
+unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="unet", torch_dtype=torch.float16)
+unet.to("cuda")
+
+# Încărcare model pentru autoencoder
+vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="vae", torch_dtype=torch.float16)
+vae.to("cuda")
+
+# Încărcare tokenizer și text model pentru CLIP
+text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+text_model.to("cuda")
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+
+# Scheduler
+scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
+
+# Pregătire optimizer
+optimizer = AdamW(unet.parameters(), lr=5e-6)
+
+# Setare model în modul de antrenament
+unet.train()
+text_model.train()
+
+# Definire număr de epoci
+num_epochs = 5  # Poți ajusta acest număr în funcție de resurse
+
+# Training loop
+for epoch in range(num_epochs):
+    for images, captions in tqdm(dataloader):
+        images = images.to("cuda", dtype=torch.float16)
+        
+        # Curăță memoria GPU înainte de fiecare iterare
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        # Tokenizare captions
+        inputs = tokenizer(captions, padding="max_length", max_length=77, return_tensors="pt").to("cuda")
+        
+        # Generare zgomot aleatoriu
+        noise = torch.randn_like(images).to("cuda", dtype=torch.float16)
+        
+        # Codificare imagini în latențe
+        latents = vae.encode(images).latent_dist.sample()
+        latents = latents * 0.18215
+
+        # Generare timesteps
+        timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (images.shape[0],), device="cuda").long()
+
+        # Forward pass prin UNet
+        encoder_hidden_states = text_model(inputs.input_ids)[0]
+        
+        # Convertim encoder_hidden_states la float16
+        encoder_hidden_states = encoder_hidden_states.to(dtype=torch.float16)
+
+        # Proiectăm dimensiunile `encoder_hidden_states` pentru a se potrivi cu cele așteptate de UNet
+        expected_dim = unet.config.cross_attention_dim
+        if encoder_hidden_states.shape[-1] != expected_dim:
+            projection_layer = torch.nn.Linear(encoder_hidden_states.shape[-1], expected_dim).to("cuda", dtype=torch.float16)
+            encoder_hidden_states = projection_layer(encoder_hidden_states)
+
+        # Generare predicție de zgomot
+        with autocast():
+            noise_pred = unet(latents, timesteps, encoder_hidden_states).sample
+
+        # Verifică dimensiunile tensorilor
+        print(f"noise_pred shape: {noise_pred.shape}")
+        print(f"noise shape: {noise.shape}")
+
+        # Redimensionare noise_pred pentru a se potrivi cu dimensiunea noise
+        if noise_pred.shape[1] != noise.shape[1]:
+            # Ajustează numărul de canale pentru noise_pred
+            conv_layer = torch.nn.Conv2d(
+                in_channels=noise_pred.shape[1],
+                out_channels=noise.shape[1],
+                kernel_size=1
+            ).to("cuda", dtype=torch.float16)
+            noise_pred = conv_layer(noise_pred)
+
+        # Redimensionare noise_pred pentru a se potrivi cu dimensiunea noise
+        if noise_pred.shape[2:] != noise.shape[2:]:
+            noise_pred = torch.nn.functional.interpolate(noise_pred, size=images.shape[2:], mode='bilinear', align_corners=False)
+
+        # Calcul pierdere (loss) comparând ieșirea modelului cu zgomotul original
+        loss = torch.nn.functional.mse_loss(noise_pred, noise)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        
+        # Curăță memoria GPU după fiecare iterare
+        gc.collect()
+        torch.cuda.empty_cache()
+    
+    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")
+
+# Salvarea modelului antrenat
+unet.save_pretrained("./finetuned-unet")
+text_model.save_pretrained("./finetuned-text-model")
+api = HfApi()
+#api.create_repo(repo_id="AndreiUrsu/finetuned-stable-diffusion-unet", repo_type="model")
+#api.create_repo(repo_id="AndreiUrsu/finetuned-stable-diffusion-text-model", repo_type="model")
+# Încărcarea pe Hugging Face
+api.upload_folder(
+    folder_path="./finetuned-unet",
+    path_in_repo=".",
+    repo_id="AndreiUrsu/finetuned-stable-diffusion-unet",
+    repo_type="model"
+)
+
+
+# Curăță memoria GPU la final
+gc.collect()
+torch.cuda.empty_cache()