Update app.py

app.py

@@ -1,67 +1,34 @@
-# Define the Pix2Pix model (UNet)
-
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from datasets import load_dataset
-from huggingface_hub import Repository, create_repo
+from huggingface_hub import Repository
 import gradio as gr
 from PIL import Image
 import os
 
-# Parameters
-IMG_SIZE = 256
-BATCH_SIZE = 1
-EPOCHS = 12
-LR = 0.0002
+from 256_model import UNet as small_UNet
+from 1024_model import UNet as big_UNet
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Define the Pix2Pix model (Simplified UNet)
-class UNet(nn.Module):
-    def __init__(self):
-        super(UNet, self).__init__()
-
-        # Encoder
-        self.encoder = nn.Sequential(
-            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),  # 256 -> 128
-            nn.ReLU(inplace=True),
-            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # 128 -> 64
-            nn.ReLU(inplace=True),
-            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # 64 -> 32
-            nn.ReLU(inplace=True),
-            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),  # 32 -> 16
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, 1024, kernel_size=4, stride=2, padding=1),  # 16 -> 8
-            nn.ReLU(inplace=True)
-        )
-
-        # Decoder
-        self.decoder = nn.Sequential(
-            nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1),  # 8 -> 16
-            nn.ReLU(inplace=True),
-            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),  # 16 -> 32
-            nn.ReLU(inplace=True),
-            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),  # 32 -> 64
-            nn.ReLU(inplace=True),
-            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # 64 -> 128
-            nn.ReLU(inplace=True),
-            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),  # 128 -> 256
-            nn.Tanh()  # Output range [-1, 1]
-        )
-
-    def forward(self, x):
-        enc = self.encoder(x)
-        dec = self.decoder(enc)
-        return dec
+big = False if device == torch.device('cpu') else True
+
+# Parameters
+IMG_SIZE = 1024 if big else 256
+BATCH_SIZE = 16 if big else 1
+EPOCHS = 12
+LR = 0.0002
+dataset_id = "K00B404/pix2pix_flux_set"
+model_repo_id = "K00B404/pix2pix_flux"
 
 # Training function
 def train_model(epochs):
     # Load the dataset
-    ds = load_dataset("K00B404/pix2pix_flux_set")
+    ds = load_dataset(dataset_id)
 
     # Transform function to resize and convert to tensor
     transform = transforms.Compose([
@@ -86,7 +53,11 @@ def train_model(epochs):
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
     # Initialize model, loss function, and optimizer
-    model = UNet().to(device)
+    try:
+        model = UNet2DModel.from_pretrained(model_repo_id).to(device)
+    except Exception:
+        model = big_UNet().to(device) if big else small_UNet().to(device)
+
     criterion = nn.L1Loss()
     optimizer = optim.Adam(model.parameters(), lr=LR)
 
@@ -97,8 +68,8 @@ def train_model(epochs):
             optimizer.zero_grad()
 
             # Forward pass
-            output = model(target)
-            loss = criterion(output, original)
+            output = model(target)  # Generate cutout image
+            loss = criterion(output, original)  # Compare with original image
 
             # Backward pass
             loss.backward()
@@ -112,21 +83,14 @@ def train_model(epochs):
 
 # Push model to Hugging Face Hub
 def push_model_to_hub(model, repo_name):
-    repo = Repository(repo_name)
-    repo.push_to_hub()
-
-    # Save the model state dict
-    model_save_path = os.path.join(repo_name, "pix2pix_model.pth")
-    torch.save(model.state_dict(), model_save_path)
-
-    # Push the model to the repo
-    repo.push_to_hub(commit_message="Initial commit with trained Pix2Pix model.")
+    # Push the model to the Hugging Face hub
+    model.push_to_hub(repo_name)
 
 # Gradio interface function
 def gradio_train(epochs):
     model = train_model(int(epochs))
-    push_model_to_hub(model, "K00B404/pix2pix_flux")
-    return f"Model trained for {epochs} epochs and pushed to Hugging Face Hub repository 'K00B404/pix2pix_flux'."
+    push_model_to_hub(model, model_repo_id)
+    return f"Model trained for {epochs} epochs on the {dataset_id} dataset and pushed to Hugging Face Hub {model_repo_id} repository."
 
 # Gradio Interface
 gr_interface = gr.Interface(
@@ -138,8 +102,9 @@ gr_interface = gr.Interface(
 )
 
 if __name__ == '__main__':
-    # Create or clone the repository
-    create_repo("K00B404/pix2pix_flux", exist_ok=True)
-    
+    # Create or clone the repository if necessary
+    repo = Repository(local_dir=model_repo_id, clone_from=model_repo_id)
+    repo.git_pull()
+
     # Launch the Gradio app
-    gr_interface.launch()
+    gr_interface.launch()