Update app.py

app.py

@@ -7,6 +7,7 @@ from datasets import load_dataset
 from huggingface_hub import Repository
 from huggingface_hub import HfApi, HfFolder, Repository, create_repo
 import os
+import pandas as pd
 import gradio as gr
 from PIL import Image
 import numpy as np
@@ -52,10 +53,17 @@ def load_model():
         return model
 
     
+import os
+import pandas as pd
 
-# Dataset class remains the same
 class Pix2PixDataset(torch.utils.data.Dataset):
-    def __init__(self, ds, transform):
+    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')
+
+        self.data = pd.read_csv(csv_path)
+        self.clip_tokenizer = clip_tokenizer
+        
         self.originals = [x for x in ds["train"] if x['label'] == 0]
         self.targets = [x for x in ds["train"] if x['label'] == 1]
         assert len(self.originals) == len(self.targets)
@@ -67,8 +75,59 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         return len(self.originals)
 
     def __getitem__(self, idx):
+        # Get original and target images
         original_img = self.originals[idx]['image']
         target_img = self.targets[idx]['image']
+        
+        # Convert PIL images
+        original = original_img.convert('RGB')
+        target = target_img.convert('RGB')
+        
+        # Extract the filename from the original image's path (assuming it has a 'filename' field or path)
+        original_img_path = self.originals[idx]['image'].filename  # Assuming it has this attribute
+        original_img_filename = os.path.basename(original_img_path)
+        
+        # Match the image filename with the `image_path` column in the CSV
+        matched_row = self.data[self.data['image_path'].str.contains(original_img_filename)]
+        
+        if matched_row.empty:
+            raise ValueError(f"No matching entry found in the CSV for image {original_img_filename}")
+
+        # Get the prompts from the matched row
+        original_prompt = matched_row['original_prompt'].values[0]
+        enhanced_prompt = matched_row['enhanced_prompt'].values[0]
+        
+        # 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 transformed images and tokenized prompts
+        return self.transform(original), self.transform(target), original_tokens, enhanced_tokens
+
+
+# Dataset class remains the same
+class Pix2PixDataset_old(torch.utils.data.Dataset):
+    def __init__(self, ds, transform, 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')
+
+        self.data = pd.read_csv(csv_path)
+        self.clip_tokenizer = clip_tokenizer
+        
+        self.originals = [x for x in ds["train"] if x['label'] == 0]
+        self.targets = [x for x in ds["train"] if x['label'] == 1]
+        assert len(self.originals) == len(self.targets)
+        print(f"Number of original images: {len(self.originals)}")
+        print(f"Number of target images: {len(self.targets)}")
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.originals)
+
+    def __getitem__(self, idx):
+        original_img = self.originals[idx]['image']
+        # TODO: get original_img file name and match with image_path in self.data....then tokenize the prompts with clip_tokenizer
+        target_img = self.targets[idx]['image']
         original = original_img.convert('RGB')
         target = target_img.convert('RGB')
         return self.transform(original), self.transform(target)
@@ -188,7 +247,7 @@ def prepare_input(image, device='cpu'):
     input_tensor = transform(image).unsqueeze(0).to(device)
     return input_tensor
 
-def run_inference(image):
+def run_inference(image, prompt):
     """Run inference on a single image"""
     global global_model
     if global_model is None:
@@ -219,41 +278,67 @@ def train_model(epochs):
         transforms.ToTensor(),
     ])
     
-    dataset = Pix2PixDataset(ds, transform)
+    # Initialize the dataset and dataloader
+    dataset = Pix2PixDataset(ds, transform, clip_tokenizer)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
     model = global_model
-    criterion = nn.L1Loss()
+    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) in enumerate(dataloader):
+        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)
-            loss = criterion(output, original)
-            loss.backward()
+
+            # 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: {loss.item():.8f}"
+                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
+    global_model = model  # Update the global model after training
     return model, "\n".join(output_text)
 
-
 def gradio_train(epochs):
     """Gradio training interface function"""
     model, training_log = train_model(int(epochs))
     to_hub(model)
     return f"{training_log}\n\nModel trained for {epochs} epochs and pushed to {model_repo_id}"
 
+def gradio_inference(input_image, keywords):
+    """Gradio inference interface function"""
+    # Generate an enhanced prompt using the chat bot
+    enhanced_prompt = chat_with_bot(keywords)
+
+    # Run inference on the input image
+    output_image = run_inference(input_image, chat_with_bot(keywords))
+
+    return input_image, output_image, keywords, enhanced_prompt
 def gradio_inference(input_image):
     """Gradio inference interface function"""
     return input_image, run_inference(input_image)