Update app.py

app.py

@@ -57,6 +57,40 @@ import os
 import pandas as pd
 
 class Pix2PixDataset(torch.utils.data.Dataset):
+    def __init__(self, combined_data, transform, clip_tokenizer):
+        self.data = combined_data
+        self.transform = transform
+        self.clip_tokenizer = clip_tokenizer
+        self.original_folder = 'images_dataset/original/'
+        self.target_folder = 'images_dataset/target/'
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        original_img_filename = os.path.basename(self.data.iloc[idx]['image_path'])
+        original_img_path = os.path.join(self.original_folder, original_img_filename)
+        target_img_path = os.path.join(self.target_folder, original_img_filename)
+
+        original_img = Image.open(original_img_path).convert('RGB')
+        target_img = Image.open(target_img_path).convert('RGB')
+        
+        # Transform images
+        original = self.transform(original_img)
+        target = self.transform(target_img)
+        
+        # Get prompts from the DataFrame
+        original_prompt = self.data.iloc[idx]['original_prompt']
+        enhanced_prompt = self.data.iloc[idx]['enhanced_prompt']
+        
+        # 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)
+        
+        return original, target, original_tokens, enhanced_tokens
+
+
+class Pix2PixDataset_older(torch.utils.data.Dataset):
     def __init__(self, ds, transform, clip_tokenizer, csv_path='combined_data.csv'):
         if not os.path.exists(csv_path):
             os.system('wget https://huggingface.co/datasets/K00B404/pix2pix_flux_set/resolve/main/combined_data.csv')
@@ -277,11 +311,73 @@ def run_inference(image, prompt):
 def to_hub(model):
     wrapper = UNetWrapper(model, model_repo_id)
     wrapper.push_to_hub()
-    
+
+
+
 def train_model(epochs):
     """Training function"""
     global global_model
     
+    # Load combined data CSV
+    data_path = 'path/to/your/combined_data.csv'  # Adjust this path
+    combined_data = pd.read_csv(data_path)
+    
+    # Define the transformation
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    # Initialize the dataset and dataloader
+    dataset = Pix2PixDataset(combined_data, transform, clip_tokenizer)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    model = global_model
+    criterion = nn.L1Loss()  # L1 loss for image reconstruction
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+    output_text = []
+    
+    for epoch in range(epochs):
+        model.train()
+        for i, (original, target, original_prompt_tokens, enhanced_prompt_tokens) in enumerate(dataloader):
+            # Move images and prompt embeddings to the appropriate device (CPU or GPU)
+            original, target = original.to(device), target.to(device)
+            original_prompt_tokens = original_prompt_tokens.input_ids.to(device)
+            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device)
+            
+            optimizer.zero_grad()
+
+            # Forward pass through the model
+            output = model(target)
+
+            # Compute image reconstruction loss
+            img_loss = criterion(output, original)
+            
+            # Compute prompt guidance loss (L2 norm between original and enhanced prompt embeddings)
+            prompt_loss = torch.norm(original_prompt_tokens - enhanced_prompt_tokens, p=2)
+            
+            # Combine losses
+            total_loss = img_loss + 0.1 * prompt_loss  # Weight the prompt guidance loss with 0.1 to balance
+            total_loss.backward()
+
+            # Optimizer step
+            optimizer.step()
+
+            if i % 10 == 0:
+                status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {total_loss.item():.8f}"
+                print(status)
+                output_text.append(status)
+        
+        # Push model to Hugging Face Hub at the end of each epoch
+        to_hub(model)
+        
+    global_model = model  # Update the global model after training
+    return model, "\n".join(output_text)
+    
+def train_model_old(epochs):
+    """Training function"""
+    global global_model
+    
     ds = load_dataset(dataset_id)
     transform = transforms.Compose([
         transforms.Resize((IMG_SIZE, IMG_SIZE)),