Create app.py

app.py

@@ -0,0 +1,145 @@
+# 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
+import gradio as gr
+from PIL import Image
+import os
+
+# Parameters
+IMG_SIZE = 256
+BATCH_SIZE = 1
+EPOCHS = 12
+LR = 0.0002
+
+# 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
+
+# Training function
+def train_model(epochs):
+    # Load the dataset
+    ds = load_dataset("K00B404/pix2pix_flux_set")
+
+    # Transform function to resize and convert to tensor
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+
+    # Create dataset and dataloader
+    class Pix2PixDataset(torch.utils.data.Dataset):
+        def __init__(self, ds):
+            self.ds = ds
+
+        def __len__(self):
+            return len(self.ds["train"])
+
+        def __getitem__(self, idx):
+            original = Image.open(self.ds["train"][idx]['original_image']).convert('RGB')
+            target = Image.open(self.ds["train"][idx]['target_image']).convert('RGB')
+            return transform(original), transform(target)
+
+    dataset = Pix2PixDataset(ds)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    # Initialize model, loss function, and optimizer
+    model = UNet().to(device)
+    criterion = nn.L1Loss()
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+
+    # Training loop
+    for epoch in range(epochs):
+        for i, (original, target) in enumerate(dataloader):
+            original, target = original.to(device), target.to(device)
+            optimizer.zero_grad()
+
+            # Forward pass
+            output = model(target)
+            loss = criterion(output, original)
+
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+
+            if i % 100 == 0:
+                print(f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.4f}")
+
+    # Return trained model
+    return model
+
+# 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.")
+
+# 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'."
+
+# Gradio Interface
+gr_interface = gr.Interface(
+    fn=gradio_train,
+    inputs=gr.Number(label="Number of Epochs"),
+    outputs="text",
+    title="Pix2Pix Model Training",
+    description="Train the Pix2Pix model and push it to the Hugging Face Hub repository."
+)
+
+if __name__ == '__main__':
+    # Create or clone the repository
+    create_repo("K00B404/pix2pix_flux", exist_ok=True)
+    
+    # Launch the Gradio app
+    gr_interface.launch()