Create app.py

app.py

@@ -0,0 +1,173 @@
+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
+from huggingface_hub import HfApi, HfFolder, Repository, create_repo
+import os
+import gradio as gr
+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
+
+# Device configuration
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+big = False if device == torch.device('cpu') else True
+
+# Parameters
+IMG_SIZE = 1024 if big else 256
+BATCH_SIZE = 16 if big else 4
+EPOCHS = 12
+LR = 0.0002
+dataset_id = "K00B404/pix2pix_flux_set"
+model_repo_id = "K00B404/pix2pix_flux"
+
+# Global model variable
+global_model = None
+
+def load_model():
+    """Load the model at startup"""
+    global global_model
+    try:
+        checkpoint = torch.load('model_weights.pth', map_location=device)
+        model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.to(device)
+        model.eval()
+        global_model = model
+        print("Model loaded successfully!")
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        model = big_UNet().to(device) if big else small_UNet().to(device)
+        global_model = model
+        return model
+
+# Dataset class remains the same
+class Pix2PixDataset(torch.utils.data.Dataset):
+    def __init__(self, ds, transform):
+        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']
+        target_img = self.targets[idx]['image']
+        original = original_img.convert('RGB')
+        target = target_img.convert('RGB')
+        return self.transform(original), self.transform(target)
+
+# UNetWrapper class remains the same
+class UNetWrapper:
+    # ... [Previous UNetWrapper implementation remains unchanged]
+    pass
+
+def prepare_input(image, device='cpu'):
+    """Prepare image for inference"""
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    input_tensor = transform(image).unsqueeze(0).to(device)
+    return input_tensor
+
+def run_inference(image):
+    """Run inference on a single image"""
+    global global_model
+    if global_model is None:
+        return "Error: Model not loaded"
+    
+    global_model.eval()
+    input_tensor = prepare_input(image, device)
+    
+    with torch.no_grad():
+        output = global_model(input_tensor)
+    
+    # Convert output to image
+    output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
+    output = ((output - output.min()) / (output.max() - output.min()) * 255).astype(np.uint8)
+    return output
+
+def train_model(epochs):
+    """Training function"""
+    global global_model
+    
+    ds = load_dataset(dataset_id)
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    dataset = Pix2PixDataset(ds, transform)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    model = global_model
+    criterion = nn.L1Loss()
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+    output_text = []
+    
+    for epoch in range(epochs):
+        model.train()
+        for i, (original, target) in enumerate(dataloader):
+            original, target = original.to(device), target.to(device)
+            optimizer.zero_grad()
+            output = model(target)
+            loss = criterion(output, original)
+            loss.backward()
+            optimizer.step()
+
+            if i % 10 == 0:
+                status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}"
+                print(status)
+                output_text.append(status)
+    
+    global_model = model
+    return model, "\n".join(output_text)
+
+def gradio_train(epochs):
+    """Gradio training interface function"""
+    model, training_log = train_model(int(epochs))
+    wrapper = UNetWrapper(model, model_repo_id)
+    wrapper.push_to_hub()
+    return f"{training_log}\n\nModel trained for {epochs} epochs and pushed to {model_repo_id}"
+
+def gradio_inference(input_image):
+    """Gradio inference interface function"""
+    return input_image, run_inference(input_image)
+
+# Create Gradio interface with tabs
+with gr.Blocks() as app:
+    gr.Markdown("# Pix2Pix Model Training and Inference")
+    
+    with gr.Tabs():
+        with gr.TabItem("Training"):
+            epochs_input = gr.Number(label="Number of Epochs")
+            train_button = gr.Button("Train Model")
+            output_text = gr.Textbox(label="Training Progress", lines=10)
+            train_button.click(gradio_train, inputs=epochs_input, outputs=output_text)
+        
+        with gr.TabItem("Inference"):
+            with gr.Row():
+                input_image = gr.Image(label="Input Image")
+                output_image = gr.Image(label="Model Output")
+            infer_button = gr.Button("Run Inference")
+            infer_button.click(gradio_inference, inputs=input_image, outputs=[input_image, output_image])
+
+if __name__ == '__main__':
+    # Load model at startup
+    load_model()
+    
+    # Launch the Gradio app
+    app.launch()