Update app.py

app.py

@@ -13,7 +13,7 @@ from PIL import Image
 import numpy as np
 from small_256_model import UNet as small_UNet
 from big_1024_model import UNet as big_UNet
-from CLIP import load as load_clip,load_vae
+from CLIP import load as load_clip,load_vae,encode_prompt
 from rich import print as rp
 
 # Device configuration
@@ -56,7 +56,7 @@ def load_model():
         return model
 
 class Pix2PixDataset(torch.utils.data.Dataset):
-    def __init__(self, combined_data, transform, clip_tokenizer):
+    def __init__(self, combined_data, transform, clip_tokenizer,clip_model):
         self.data = combined_data
         self.transform = transform
         self.clip_tokenizer = clip_tokenizer
@@ -81,12 +81,24 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         # Get prompts from the DataFrame
         original_prompt = self.data.iloc[idx]['original_prompt']
         enhanced_prompt = self.data.iloc[idx]['enhanced_prompt']
+
+        # Encode images
+        original_image_latents = vae.encode(original_images).latent_dist.sample()
+        target_image_latents = vae.encode(target_images).latent_dist.sample()
         
+        # Encode prompts
+        prompt_latents = encode_prompt(enhanced_prompt,clip_model,clip_tokenizer)
+
+        # Pass these to your Pix2Pix model
+        #generated_images = pix2pix_model(original_latents, prompt_latents)
+
+
+        return original_image_latents,target_image_latents,prompt_latents
         # Tokenize the prompts using CLIP tokenizer
-        original_tokens = self.clip_tokenizer(original_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
-        enhanced_tokens = self.clip_tokenizer(enhanced_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
+        #original_tokens = self.clip_tokenizer(original_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
+        #enhanced_tokens = self.clip_tokenizer(enhanced_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
         
-        return original, target, original_tokens, enhanced_tokens
+        #return original, target, original_tokens, enhanced_tokens
 
 
 
@@ -261,11 +273,11 @@ def train_model(epochs, save_interval=1):
     ])
     
     # Initialize dataset and dataloader
-    dataset = Pix2PixDataset(combined_data, transform, clip_tokenizer)
+    dataset = Pix2PixDataset(combined_data, transform, clip_tokenizer, clip_model)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
     model = global_model
-    criterion = nn.L1Loss()
+    criterion = nn.L1Loss()  # You may change this to suit your loss calculation needs
     optimizer = optim.Adam(model.parameters(), lr=LR)
     scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)  # Example scheduler
     wrapper = UNetWrapper(model, model_repo_id, epoch=0, loss=0.0, optimizer=optimizer, scheduler=scheduler)
@@ -276,17 +288,17 @@ def train_model(epochs, save_interval=1):
         model.train()
         running_loss = 0.0
 
-        for i, (original, target, original_prompt_tokens, enhanced_prompt_tokens) in enumerate(dataloader):
+        for i, (latent_original, latent_target, latent_prompt) in enumerate(dataloader):
             # Move data to device
-            original, target = original.to(device), target.to(device)
-            original_prompt_tokens = original_prompt_tokens.input_ids.to(device).float()
-            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device).float()
-            
+            latent_original, latent_target, latent_prompt = latent_original.to(device), latent_target.to(device), latent_prompt.to(device)
+
             optimizer.zero_grad()
 
-            # Forward pass
-            output = model(target)
-            img_loss = criterion(output, original)
+            # Forward pass with the latents
+            output = model(latent_target, latent_prompt)  # Assuming your model can take both target and prompt latents
+
+            # Calculate loss using the original latents
+            img_loss = criterion(output, latent_original)
             total_loss = img_loss
             total_loss.backward()
             optimizer.step()
@@ -304,7 +316,7 @@ def train_model(epochs, save_interval=1):
         
         # Save checkpoint at specified intervals
         if (epoch + 1) % save_interval == 0:
-            checkpoint_path = f'big_checkpoint_epoch_{epoch+1}.pth' if big else   f'small_checkpoint_epoch_{epoch+1}.pth'
+            checkpoint_path = f'big_checkpoint_epoch_{epoch+1}.pth' if big else f'small_checkpoint_epoch_{epoch+1}.pth'
             wrapper.save_checkpoint(checkpoint_path)
             wrapper.push_to_hub(checkpoint_path)