huggingface spaces commit.

.gitattributes

@@ -1,35 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+*.pyc
+.env
+__pycache__
+*.pth.tar
+results/
+mlruns/
+# checkpoints/

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.8
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY . .
+
+CMD ["python", "data_ingestion.py"]

README.md

@@ -1,12 +1,11 @@
----
-title: Pix2pixcolorizer
-emoji: 📊
-colorFrom: gray
-colorTo: blue
-sdk: gradio
-sdk_version: 4.42.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Image Colorizer
+
+This is a Gradio app that colorizes grayscale images using a trained GAN model. Upload a grayscale image, and the app will return a colorized version of the image.
+
+## How to use
+
+1. Upload a grayscale image using the provided interface.
+2. Wait for the model to process the image.
+3. View the colorized result!
+
+This app uses a GAN model trained on the ImageNet dataset to add color to grayscale images.

app-mlflow.py

@@ -0,0 +1,72 @@
+import gradio as gr
+import torch
+import mlflow
+import numpy as np
+from PIL import Image
+from skimage.color import rgb2lab, lab2rgb
+from torchvision import transforms
+
+from model import Generator
+
+EXPERIMENT_NAME = "Colorizer_Experiment"
+RUN_ID = "your_run_id_here"  # Replace with your actual run ID
+
+def setup_mlflow():
+    experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
+    else:
+        experiment_id = experiment.experiment_id
+    return experiment_id
+
+def load_model(run_id, device):
+    print(f"Loading model from run: {run_id}")
+    model_uri = f"runs:/{run_id}/generator_model"
+    model = mlflow.pytorch.load_model(model_uri, map_location=device)
+    return model
+
+def preprocess_image(image):
+    img = Image.fromarray(image).convert("RGB")
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor()
+    ])
+    img_tensor = transform(img)
+    lab_img = rgb2lab(img_tensor.permute(1, 2, 0).numpy())
+    L = lab_img[:,:,0]
+    L = (L - 50) / 50
+    L = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float()
+    return L
+
+def postprocess_output(L, ab):
+    L = L.squeeze().cpu().numpy()
+    ab = ab.squeeze().cpu().numpy()
+    L = (L + 1.) * 50.
+    ab = ab * 128.
+    Lab = np.concatenate([L[..., np.newaxis], ab], axis=2)
+    rgb_img = lab2rgb(Lab)
+    return (rgb_img * 255).astype(np.uint8)
+
+def colorize_image(image, model, device):
+    L = preprocess_image(image).to(device)
+    with torch.no_grad():
+        ab = model(L)
+    colorized = postprocess_output(L, ab)
+    return colorized
+
+def setup_gradio_app(run_id, device):
+    model = load_model(run_id, device)
+
+    def gradio_colorize(input_image):
+        colorized = colorize_image(input_image, model, device)
+        return Image.fromarray(colorized)
+
+    iface = gr.Interface(
+        fn=gradio_colorize,
+        inputs=gr.Image(label="Upload a grayscale image"),
+        outputs=gr.Image(label="Colorized Image"),
+        title="Image Colorizer",
+        description="Upload a grayscale image and get a colorized version!",
+    )
+
+    return iface

app.py

@@ -0,0 +1,123 @@
+import gradio as gr
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+from skimage.color import rgb2lab, lab2rgb
+import os
+from torch import nn
+
+# Define the model architecture (same as in the training script)
+class UNetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, down=True, bn=True, dropout=False):
+        super(UNetBlock, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False) if down \
+            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels) if bn else None
+        self.dropout = nn.Dropout(0.5) if dropout else None
+        self.down = down
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn:
+            x = self.bn(x)
+        if self.dropout:
+            x = self.dropout(x)
+        return nn.ReLU()(x) if self.down else nn.ReLU(inplace=True)(x)
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.down1 = UNetBlock(1, 64, bn=False)
+        self.down2 = UNetBlock(64, 128)
+        self.down3 = UNetBlock(128, 256)
+        self.down4 = UNetBlock(256, 512)
+        self.down5 = UNetBlock(512, 512)
+        self.down6 = UNetBlock(512, 512)
+        self.down7 = UNetBlock(512, 512)
+        self.down8 = UNetBlock(512, 512, bn=False)
+        
+        self.up1 = UNetBlock(512, 512, down=False, dropout=True)
+        self.up2 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up3 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up4 = UNetBlock(1024, 512, down=False)
+        self.up5 = UNetBlock(1024, 256, down=False)
+        self.up6 = UNetBlock(512, 128, down=False)
+        self.up7 = UNetBlock(256, 64, down=False)
+        self.up8 = nn.ConvTranspose2d(128, 2, 4, 2, 1)
+
+    def forward(self, x):
+        d1 = self.down1(x)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+        d6 = self.down6(d5)
+        d7 = self.down7(d6)
+        d8 = self.down8(d7)
+        
+        u1 = self.up1(d8)
+        u2 = self.up2(torch.cat([u1, d7], 1))
+        u3 = self.up3(torch.cat([u2, d6], 1))
+        u4 = self.up4(torch.cat([u3, d5], 1))
+        u5 = self.up5(torch.cat([u4, d4], 1))
+        u6 = self.up6(torch.cat([u5, d3], 1))
+        u7 = self.up7(torch.cat([u6, d2], 1))
+        return torch.tanh(self.up8(torch.cat([u7, d1], 1)))
+
+# Load the checkpoint
+def load_checkpoint(filename, generator, map_location):
+    if os.path.isfile(filename):
+        print(f"Loading checkpoint '{filename}'")
+        checkpoint = torch.load(filename, map_location=map_location)
+        generator.load_state_dict(checkpoint['generator_state_dict'])
+        print(f"Loaded checkpoint '{filename}' (epoch {checkpoint['epoch']})")
+    else:
+        print(f"No checkpoint found at '{filename}'")
+
+# Initialize the model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+generator = Generator().to(device)
+checkpoint_path = "checkpoints/latest_checkpoint.pth.tar"
+load_checkpoint(checkpoint_path, generator, map_location=device)
+generator.eval()
+
+# Define the transformation
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.Grayscale(num_output_channels=1),  # Convert to grayscale
+    transforms.ToTensor()
+])
+
+# Define the inference function
+def colorize_image(input_image):
+    try:
+        original_size = input_image.size
+        input_image = transform(input_image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            output = generator(input_image)
+        output = output.squeeze(0).cpu().numpy()
+        L = input_image.squeeze(0).cpu().numpy()
+        L = (L + 1.) * 50.
+        ab = output * 128.
+        Lab = np.concatenate([L, ab], axis=0).transpose(1, 2, 0)
+        rgb_image = lab2rgb(Lab)
+        rgb_image = Image.fromarray((rgb_image * 255).astype(np.uint8))
+        rgb_image = rgb_image.resize(original_size, Image.LANCZOS)
+        return rgb_image
+    except Exception as e:
+        print(f"Error in colorize_image: {str(e)}")
+        return None
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=colorize_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Image Colorizer",
+    description="Upload a grayscale image to colorize it."
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()

colorizer_pipeline.py

@@ -0,0 +1,379 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import IterableDataset, DataLoader
+from torchvision import transforms
+from torchvision.utils import make_grid
+import mlflow
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import numpy as np
+from skimage.color import rgb2lab, lab2rgb
+from datasets import load_dataset
+from PIL import Image
+from itertools import islice
+import traceback
+
+# MLflow setup
+EXPERIMENT_NAME = "Colorizer_Experiment"
+
+def setup_mlflow():
+    experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
+    else:
+        experiment_id = experiment.experiment_id
+    return experiment_id
+
+# Data ingestion
+class ColorizeIterableDataset(IterableDataset):
+    def __init__(self, dataset, transform=None):
+        self.dataset = dataset
+        self.transform = transform
+
+    def __iter__(self):
+        for item in self.dataset:
+            try:
+                img = item['image']
+                if img.mode != 'RGB':
+                    img = img.convert('RGB')
+                if self.transform:
+                    img = self.transform(img)
+                
+                # Add shape check after transform
+                if img.shape != (3, 256, 256):
+                    print(f"Unexpected image shape after transform: {img.shape}")
+                    continue
+
+                lab = rgb2lab(img.permute(1, 2, 0).numpy())
+                
+                # Add shape check after rgb2lab conversion
+                if lab.shape != (256, 256, 3):
+                    print(f"Unexpected lab shape: {lab.shape}")
+                    continue
+
+                l_chan = lab[:, :, 0]
+                l_chan = (l_chan - 50) / 50
+                ab_chan = lab[:, :, 1:]
+                ab_chan = ab_chan / 128
+                
+                yield torch.Tensor(l_chan).unsqueeze(0), torch.Tensor(ab_chan).permute(2, 0, 1)
+            except Exception as e:
+                print(f"Error processing image: {str(e)}")
+                continue
+
+def create_dataloaders(batch_size=32):
+    try:
+        print("Loading ImageNet dataset in streaming mode...")
+        dataset = load_dataset("imagenet-1k", split="train", streaming=True)
+        print("Dataset loaded in streaming mode.")
+        
+        print("Creating custom dataset...")
+        transform = transforms.Compose([
+            transforms.Resize((256, 256)),  # Resize all images to 256x256
+            transforms.ToTensor()
+        ])
+        train_dataset = ColorizeIterableDataset(dataset, transform=transform)
+        print("Custom dataset created.")
+        
+        print("Creating dataloader...")
+        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=4)
+        print("Dataloader created.")
+        
+        return train_dataloader
+    except Exception as e:
+        print(f"Error in create_dataloaders: {str(e)}")
+        return None
+
+# Model definition
+class UNetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, down=True, bn=True, dropout=False):
+        super(UNetBlock, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False) if down \
+            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels) if bn else None
+        self.dropout = nn.Dropout(0.5) if dropout else None
+        self.down = down
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn:
+            x = self.bn(x)
+        if self.dropout:
+            x = self.dropout(x)
+        return nn.ReLU()(x) if self.down else nn.ReLU(inplace=True)(x)
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.down1 = UNetBlock(1, 64, bn=False)
+        self.down2 = UNetBlock(64, 128)
+        self.down3 = UNetBlock(128, 256)
+        self.down4 = UNetBlock(256, 512)
+        self.down5 = UNetBlock(512, 512)
+        self.down6 = UNetBlock(512, 512)
+        self.down7 = UNetBlock(512, 512)
+        self.down8 = UNetBlock(512, 512, bn=False)
+        
+        self.up1 = UNetBlock(512, 512, down=False, dropout=True)
+        self.up2 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up3 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up4 = UNetBlock(1024, 512, down=False)
+        self.up5 = UNetBlock(1024, 256, down=False)
+        self.up6 = UNetBlock(512, 128, down=False)
+        self.up7 = UNetBlock(256, 64, down=False)
+        self.up8 = nn.ConvTranspose2d(128, 2, 4, 2, 1)
+
+    def forward(self, x):
+        d1 = self.down1(x)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+        d6 = self.down6(d5)
+        d7 = self.down7(d6)
+        d8 = self.down8(d7)
+        
+        u1 = self.up1(d8)
+        u2 = self.up2(torch.cat([u1, d7], 1))
+        u3 = self.up3(torch.cat([u2, d6], 1))
+        u4 = self.up4(torch.cat([u3, d5], 1))
+        u5 = self.up5(torch.cat([u4, d4], 1))
+        u6 = self.up6(torch.cat([u5, d3], 1))
+        u7 = self.up7(torch.cat([u6, d2], 1))
+        return torch.tanh(self.up8(torch.cat([u7, d1], 1)))
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(3, 64, 4, stride=2, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(256, 512, 4, padding=1),
+            nn.BatchNorm2d(512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(512, 1, 4, padding=1)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+def init_weights(model):
+    classname = model.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(model.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(model.weight.data, 1.0, 0.02)
+        nn.init.constant_(model.bias.data, 0)
+
+# Training utilities
+def lab_to_rgb(L, ab):
+    L = (L + 1.) * 50.
+    ab = ab * 128.
+    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
+    rgb_imgs = []
+    for img in Lab:
+        img_rgb = lab2rgb(img)
+        rgb_imgs.append(img_rgb)
+    return np.stack(rgb_imgs, axis=0)
+
+def visualize_results(epoch, generator, train_loader, device):
+    generator.eval()
+    with torch.no_grad():
+        try:
+            for inputs, real_AB in train_loader:
+                print(f"Input shape: {inputs.shape}, real_AB shape: {real_AB.shape}")
+                
+                # Ensure inputs have the correct shape (B, 1, H, W)
+                if inputs.shape[1] != 1:
+                    inputs = inputs.unsqueeze(1)
+                
+                inputs, real_AB = inputs.to(device), real_AB.to(device)
+                fake_AB = generator(inputs)
+                
+                print(f"fake_AB shape: {fake_AB.shape}")
+                
+                # Ensure fake_AB and real_AB have the correct shape (B, 2, H, W)
+                if fake_AB.shape[1] != 2:
+                    fake_AB = fake_AB.view(fake_AB.shape[0], 2, fake_AB.shape[2], fake_AB.shape[3])
+                if real_AB.shape[1] != 2:
+                    real_AB = real_AB.view(real_AB.shape[0], 2, real_AB.shape[2], real_AB.shape[3])
+                
+                fake_rgb = lab_to_rgb(inputs.cpu(), fake_AB.cpu())
+                real_rgb = lab_to_rgb(inputs.cpu(), real_AB.cpu())
+                
+                print(f"fake_rgb shape: {fake_rgb.shape}, real_rgb shape: {real_rgb.shape}")
+                
+                concatenated = np.concatenate([real_rgb, fake_rgb], axis=2)  # Changed axis from 3 to 2
+                print(f"Concatenated shape: {concatenated.shape}")
+                
+                img_grid = make_grid(torch.from_numpy(concatenated).permute(0, 3, 1, 2), normalize=True, nrow=4)
+                
+                plt.figure(figsize=(15, 15))
+                plt.imshow(img_grid.permute(1, 2, 0).cpu())
+                plt.axis('off')
+                plt.title(f'Epoch {epoch}')
+                plt.savefig(f'results/epoch_{epoch}.png')
+                mlflow.log_artifact(f'results/epoch_{epoch}.png')
+                plt.close()
+                break
+        except Exception as e:
+            print(f"Error in visualize_results: {str(e)}")
+            traceback.print_exc()
+    generator.train()
+
+def save_checkpoint(state, filename="checkpoint.pth.tar"):
+    torch.save(state, filename)
+    mlflow.log_artifact(filename)
+
+def load_checkpoint(filename, generator, discriminator, optimizerG, optimizerD):
+    if os.path.isfile(filename):
+        print(f"Loading checkpoint '{filename}'")
+        checkpoint = torch.load(filename)
+        start_epoch = checkpoint['epoch'] + 1
+        generator.load_state_dict(checkpoint['generator_state_dict'])
+        discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
+        optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])
+        optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])
+        print(f"Loaded checkpoint '{filename}' (epoch {checkpoint['epoch']})")
+        return start_epoch
+    else:
+        print(f"No checkpoint found at '{filename}'")
+        return 0
+
+# Training function
+def train(generator, discriminator, train_loader, num_epochs, device, lr=0.0002, beta1=0.5):
+    criterion = nn.BCEWithLogitsLoss()
+    l1_loss = nn.L1Loss()
+
+    optimizerG = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
+    optimizerD = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))
+
+    checkpoint_dir = "checkpoints"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    os.makedirs("results", exist_ok=True)
+
+    checkpoint_path = os.path.join(checkpoint_dir, "latest_checkpoint.pth.tar")
+    start_epoch = load_checkpoint(checkpoint_path, generator, discriminator, optimizerG, optimizerD)
+
+    experiment_id = setup_mlflow()
+    with mlflow.start_run(experiment_id=experiment_id, run_name="training_run") as run:
+        try:
+            for epoch in range(start_epoch, num_epochs):
+                generator.train()
+                discriminator.train()
+
+                num_iterations = 2
+                pbar = tqdm(enumerate(islice(train_loader, num_iterations)), total=num_iterations, desc=f"Epoch {epoch+1}/{num_epochs}")
+                
+                for i, (real_L, real_AB) in pbar:
+                    # Add shape check
+                    if real_L.shape[1:] != (1, 256, 256) or real_AB.shape[1:] != (2, 256, 256):
+                        print(f"Unexpected tensor shapes: real_L {real_L.shape}, real_AB {real_AB.shape}")
+                        continue
+
+                    real_L, real_AB = real_L.to(device), real_AB.to(device)
+                    batch_size = real_L.size(0)
+
+                    # Train Discriminator
+                    optimizerD.zero_grad()
+
+                    fake_AB = generator(real_L)
+                    fake_LAB = torch.cat([real_L, fake_AB], dim=1)
+                    real_LAB = torch.cat([real_L, real_AB], dim=1)
+
+                    pred_fake = discriminator(fake_LAB.detach())
+                    loss_D_fake = criterion(pred_fake, torch.zeros_like(pred_fake))
+
+                    pred_real = discriminator(real_LAB)
+                    loss_D_real = criterion(pred_real, torch.ones_like(pred_real))
+
+                    loss_D = (loss_D_fake + loss_D_real) * 0.5
+                    loss_D.backward()
+                    optimizerD.step()
+
+                    # Train Generator
+                    optimizerG.zero_grad()
+
+                    fake_AB = generator(real_L)
+                    fake_LAB = torch.cat([real_L, fake_AB], dim=1)
+                    pred_fake = discriminator(fake_LAB)
+
+                    loss_G_GAN = criterion(pred_fake, torch.ones_like(pred_fake))
+                    loss_G_L1 = l1_loss(fake_AB, real_AB) * 100  # L1 loss weight
+
+                    loss_G = loss_G_GAN + loss_G_L1
+                    loss_G.backward()
+                    optimizerG.step()
+
+                    pbar.set_postfix({
+                        'D_loss': loss_D.item(),
+                        'G_loss': loss_G.item(),
+                        'G_L1': loss_G_L1.item()
+                    })  
+
+                    mlflow.log_metrics({
+                        "D_loss": loss_D.item(),
+                        "G_loss": loss_G.item(),
+                        "G_L1_loss": loss_G_L1.item()
+                    }, step=epoch * num_iterations + i)
+
+                visualize_results(epoch, generator, train_loader, device)
+
+                checkpoint = {
+                    'epoch': epoch,
+                    'generator_state_dict': generator.state_dict(),
+                    'discriminator_state_dict': discriminator.state_dict(),
+                    'optimizerG_state_dict': optimizerG.state_dict(),
+                    'optimizerD_state_dict': optimizerD.state_dict(),
+                }
+                save_checkpoint(checkpoint, filename=checkpoint_path)
+
+            print("Training completed successfully.")
+            
+            mlflow.pytorch.log_model(generator, "generator_model")
+            model_uri = f"runs:/{run.info.run_id}/generator_model"
+            mlflow.register_model(model_uri, "colorizer_generator")
+            
+            return run.info.run_id
+            
+        except Exception as e:
+            print(f"Error during training: {str(e)}")
+            mlflow.log_param("error", str(e))
+            return None
+
+# Main execution
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    try:
+        batch_size = 32
+        num_epochs = 50
+
+        train_loader = create_dataloaders(batch_size=batch_size)
+        if train_loader is None:
+            raise Exception("Failed to create dataloader")
+
+        generator = Generator().to(device)
+        discriminator = Discriminator().to(device)
+
+        generator.apply(init_weights)
+        discriminator.apply(init_weights)
+
+        run_id = train(generator, discriminator, train_loader, num_epochs=num_epochs, device=device)
+        if run_id:
+            print(f"Training completed successfully. Run ID: {run_id}")
+            with open("latest_run_id.txt", "w") as f:
+                f.write(run_id)
+        else:
+            print("Training failed!")
+    except Exception as e:
+        print(f"Critical error in main execution: {str(e)}")

data_ingestion.py

@@ -0,0 +1,138 @@
+import os
+import mlflow
+import torch
+from torch.utils.data import IterableDataset, DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+from skimage.color import rgb2lab
+from PIL import Image
+import numpy as np
+
+
+class ColorizeIterableDataset(IterableDataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.transform = transforms.Compose([
+            transforms.Resize((256, 256)),
+            transforms.ToTensor()
+        ])
+
+    def __iter__(self):
+        for item in self.dataset:
+            try:
+                img = item['image']
+                if img.mode != 'RGB':
+                    img = img.convert('RGB')
+                img = self.transform(img)
+                
+                # Convert to LAB color space
+                lab = rgb2lab(img.permute(1, 2, 0).numpy())
+                
+                # Normalize L channel to range [-1, 1]
+                l_chan = lab[:, :, 0]
+                l_chan = (l_chan - 50) / 50
+                
+                # Normalize AB channels to range [-1, 1]
+                ab_chan = lab[:, :, 1:]
+                ab_chan = ab_chan / 128
+                
+                yield torch.Tensor(l_chan).unsqueeze(0), torch.Tensor(ab_chan).permute(2, 0, 1)
+            except Exception as e:
+                print(f"Error processing image: {str(e)}")
+                continue
+
+def create_dataloaders(batch_size=32):
+    try:
+        print("Loading ImageNet dataset in streaming mode...")
+        # Load ImageNet dataset from Hugging Face in streaming mode
+        dataset = load_dataset("imagenet-1k", split="train", streaming=True)
+        print("Dataset loaded in streaming mode.")
+        
+        print("Creating custom dataset...")
+        # Create custom dataset
+        train_dataset = ColorizeIterableDataset(dataset)
+        print("Custom dataset created.")
+        
+        print("Creating dataloader...")
+        # Create dataloader
+        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=4)
+        print("Dataloader created.")
+        
+        return train_dataloader
+    except Exception as e:
+        print(f"Error in create_dataloaders: {str(e)}")
+        return None
+
+def test_data_ingestion():
+    print("Testing data ingestion...")
+    try:
+        dataloader = create_dataloaders(batch_size=4)
+        if dataloader is None:
+            raise Exception("Dataloader creation failed")
+        
+        # Get the first batch
+        for sample_batch in dataloader:
+            if len(sample_batch) != 2:
+                raise Exception(f"Unexpected batch format: {len(sample_batch)} elements instead of 2")
+            
+            l_chan, ab_chan = sample_batch
+            if l_chan.shape != torch.Size([4, 1, 256, 256]) or ab_chan.shape != torch.Size([4, 2, 256, 256]):
+                raise Exception(f"Unexpected tensor shapes: L={l_chan.shape}, AB={ab_chan.shape}")
+            
+            print("Data ingestion test passed.")
+            return True
+    except Exception as e:
+        print(f"Data ingestion test failed: {str(e)}")
+        return False
+
+if __name__ == "__main__":
+    try:
+        print("Starting data ingestion pipeline...")
+        mlflow.start_run(run_name="data_ingestion")
+        
+        try:
+            # Log parameters
+            print("Logging parameters...")
+            mlflow.log_param("batch_size", 32)
+            mlflow.log_param("dataset", "imagenet-1k")
+            print("Parameters logged.")
+            
+            # Create dataloaders
+            print("Creating dataloaders...")
+            train_dataloader = create_dataloaders(batch_size=32)
+            if train_dataloader is None:
+                raise Exception("Failed to create dataloader")
+            print("Dataloaders created successfully.")
+            
+            # Log a sample batch
+            print("Logging sample batch...")
+            for sample_batch in train_dataloader:
+                l_chan, ab_chan = sample_batch
+                
+                # Log sample input (L channel)
+                sample_input = l_chan[0].numpy()
+                mlflow.log_image(sample_input, "sample_input_l_channel.png")
+                
+                # Log sample target (AB channels)
+                sample_target = ab_chan[0].permute(1, 2, 0).numpy()
+                mlflow.log_image(sample_target, "sample_target_ab_channels.png")
+                
+                print("Sample batch logged.")
+                break  # We only need one batch for logging
+            
+            print("Data ingestion pipeline completed successfully.")
+        
+        except Exception as e:
+            print(f"Error in data ingestion pipeline: {str(e)}")
+            mlflow.log_param("error", str(e))
+        
+        finally:
+            mlflow.end_run()
+    
+    except Exception as e:
+        print(f"Critical error in main execution: {str(e)}")
+
+    if test_data_ingestion():
+        print("All tests passed.")
+    else:
+        print("Tests failed.")

inference.py

@@ -0,0 +1,117 @@
+import os
+import torch
+import mlflow
+import mlflow.pytorch
+from PIL import Image
+import numpy as np
+from skimage.color import rgb2lab, lab2rgb
+from torchvision import transforms
+import argparse
+
+from model import Generator
+
+EXPERIMENT_NAME = "Colorizer_Experiment"
+
+def setup_mlflow():
+    experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
+    else:
+        experiment_id = experiment.experiment_id
+    return experiment_id
+
+def load_model(run_id, device):
+    print(f"Loading model from run: {run_id}")
+    model_uri = f"runs:/{run_id}/generator_model"
+    model = mlflow.pytorch.load_model(model_uri, map_location=device)
+    return model
+
+# Configuration variables
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+RUN_ID = "your_run_id_here"  # Replace with the actual run ID
+IMAGE_PATH = "path/to/your/image.jpg"  # Replace with the path to your input image
+SAVE_MODEL = False
+SERVE_MODEL = False
+SERVE_PORT = 5000
+
+def preprocess_image(image_path):
+    img = Image.open(image_path).convert("RGB")
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor()
+    ])
+    img_tensor = transform(img)
+    lab_img = rgb2lab(img_tensor.permute(1, 2, 0).numpy())
+    L = lab_img[:,:,0]
+    L = (L - 50) / 50
+    L = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float()
+    return L
+
+def postprocess_output(L, ab):
+    L = L.squeeze().cpu().numpy()
+    ab = ab.squeeze().cpu().numpy()
+    L = (L + 1.) * 50.
+    ab = ab * 128.
+    Lab = np.concatenate([L[..., np.newaxis], ab], axis=2)
+    rgb_img = lab2rgb(Lab)
+    return (rgb_img * 255).astype(np.uint8)
+
+def colorize_image(model, image_path, device):
+    L = preprocess_image(image_path).to(device)
+    with torch.no_grad():
+        ab = model(L)
+    colorized = postprocess_output(L, ab)
+    return colorized
+
+def save_model(model, run_id):
+    with mlflow.start_run(run_id=run_id):
+        # Log the model
+        mlflow.pytorch.log_model(model, "model")
+        
+        # Register the model
+        model_uri = f"runs:/{run_id}/model"
+        mlflow.register_model(model_uri, "colorizer_model")
+        
+        print(f"Model saved and registered with run_id: {run_id}")
+
+def serve_model(run_id, port=5000):
+    model_uri = f"runs:/{run_id}/model"
+    mlflow.pytorch.serve(model_uri, port=port)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Colorizer Inference")
+    parser.add_argument("--run_id", type=str, help="MLflow run ID of the trained model")
+    parser.add_argument("--image_path", type=str, required=True, help="Path to the input grayscale image")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
+                        help="Device to use for inference (cuda/cpu)")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    # If run_id is not provided, try to load it from the file
+    if not args.run_id:
+        try:
+            with open("latest_run_id.txt", "r") as f:
+                args.run_id = f.read().strip()
+        except FileNotFoundError:
+            print("No run ID provided and couldn't find latest_run_id.txt")
+            exit(1)
+
+    experiment_id = setup_mlflow()
+    with mlflow.start_run(experiment_id=experiment_id, run_name="inference_run"):
+        try:
+            model = load_model(args.run_id, device)
+            colorized = colorize_image(model, args.image_path, device)
+            output_path = f"colorized_{os.path.basename(args.image_path)}"
+            Image.fromarray(colorized).save(output_path)
+            print(f"Colorized image saved as: {output_path}")
+            
+            mlflow.log_artifact(output_path)
+            mlflow.log_param("input_image", args.image_path)
+            mlflow.log_param("model_run_id", args.run_id)
+        except Exception as e:
+            print(f"Error during inference: {str(e)}")
+            mlflow.log_param("error", str(e))
+        finally:
+            mlflow.end_run()

instructions.txt

@@ -0,0 +1,13 @@
+Create a Colorizer webapp. This is a professional project and should follow all necessary coding practices, machine learning as well as MLops practices. Use python, mlflow, etc as well as other libraries and tools of your choice which might be required to achieve a fully functioning app. Can deploy it on gradio.
+1. Gather a large and diverse dataset of color images. Convert these images to grayscale to create the paried training dataset for the Pix2Pix model.
+2. Implement a Pix2Pix GAN architecture, utilizing U net generator for colorization and CNN discriminator for realism evaluation. Ensure the model processes images in the LAB color space.
+3. Define the loss functions for training: l1 loss for pixel accuracy and adversarial loss for generating realistic images.
+4. Train the Pix2Pix model using grayscale images as input and color images as targets. Monitor validation performance and apply data augmentation techniques.
+5. Evaluate the trained model on the test dataset using PSNR and SSIM metrics. generate and visually inspect colorized images.
+5. Develop a user interface for application that allows users to upload images. use gradio for this. Integrate the trained model and ensure the output maintains the original image size and quality.
+6. Add post processing features to enhance the quality of colorized images. implement a feedback mechanism for users.
+7. Perform user testing to gather feedback on application. Use this feedback to iterate and improve the model and user interface.
+
+This is a professional project and should follow all necessary coding practices, machine learning as well as MLops practices. Use python, mlflow, etc as well as other libraries and tools of your choice which might be required to achieve a fully functioning app.
+All the functions should be completely implemented and working. 
+The project should be containerized and should be able to run in a docker container. 

model.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class UNetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, down=True, bn=True, dropout=False):
+        super(UNetBlock, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False) if down \
+            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels) if bn else None
+        self.dropout = nn.Dropout(0.5) if dropout else None
+        self.down = down
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn:
+            x = self.bn(x)
+        if self.dropout:
+            x = self.dropout(x)
+        return F.relu(x) if self.down else F.relu(x, inplace=True)
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.down1 = UNetBlock(1, 64, bn=False)  # Input is L channel (1 channel)
+        self.down2 = UNetBlock(64, 128)
+        self.down3 = UNetBlock(128, 256)
+        self.down4 = UNetBlock(256, 512)
+        self.down5 = UNetBlock(512, 512)
+        self.down6 = UNetBlock(512, 512)
+        self.down7 = UNetBlock(512, 512)
+        self.down8 = UNetBlock(512, 512, bn=False)
+        
+        self.up1 = UNetBlock(512, 512, down=False, dropout=True)
+        self.up2 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up3 = UNetBlock(1024, 512, down=False, dropout=True)
+        self.up4 = UNetBlock(1024, 512, down=False)
+        self.up5 = UNetBlock(1024, 256, down=False)
+        self.up6 = UNetBlock(512, 128, down=False)
+        self.up7 = UNetBlock(256, 64, down=False)
+        self.up8 = nn.ConvTranspose2d(128, 2, 4, 2, 1)  # Output is AB channels (2 channels)
+
+    def forward(self, x):
+        d1 = self.down1(x)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+        d6 = self.down6(d5)
+        d7 = self.down7(d6)
+        d8 = self.down8(d7)
+        
+        u1 = self.up1(d8)
+        u2 = self.up2(torch.cat([u1, d7], 1))
+        u3 = self.up3(torch.cat([u2, d6], 1))
+        u4 = self.up4(torch.cat([u3, d5], 1))
+        u5 = self.up5(torch.cat([u4, d4], 1))
+        u6 = self.up6(torch.cat([u5, d3], 1))
+        u7 = self.up7(torch.cat([u6, d2], 1))
+        return torch.tanh(self.up8(torch.cat([u7, d1], 1)))
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(3, 64, 4, stride=2, padding=1),  # Input is L+AB (3 channels)
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(256, 512, 4, padding=1),
+            nn.BatchNorm2d(512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(512, 1, 4, padding=1)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+def init_weights(model):
+    classname = model.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(model.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(model.weight.data, 1.0, 0.02)
+        nn.init.constant_(model.bias.data, 0)
+
+def create_models():
+    try:
+        print("Creating Generator...")
+        generator = Generator()
+        generator.apply(init_weights)
+        print("Generator created successfully.")
+
+        print("Creating Discriminator...")
+        discriminator = Discriminator()
+        discriminator.apply(init_weights)
+        print("Discriminator created successfully.")
+
+        return generator, discriminator
+    except Exception as e:
+        print(f"Error in creating models: {str(e)}")
+        return None, None
+
+def test_models():
+    print("Testing models...")
+    try:
+        generator, discriminator = create_models()
+        if generator is None or discriminator is None:
+            raise Exception("Model creation failed")
+        
+        test_input_g = torch.randn(1, 1, 256, 256)
+        test_output_g = generator(test_input_g)
+        if test_output_g.shape != torch.Size([1, 2, 256, 256]):
+            raise Exception(f"Unexpected generator output shape: {test_output_g.shape}")
+        
+        test_input_d = torch.randn(1, 3, 256, 256)
+        test_output_d = discriminator(test_input_d)
+        if test_output_d.shape != torch.Size([1, 1, 30, 30]):
+            raise Exception(f"Unexpected discriminator output shape: {test_output_d.shape}")
+        
+        print("Model test passed.")
+        return True
+    except Exception as e:
+        print(f"Model test failed: {str(e)}")
+        return False
+
+if __name__ == "__main__":
+    try:
+        print("Initializing models...")
+        generator, discriminator = create_models()
+
+        if generator is None or discriminator is None:
+            raise Exception("Failed to create models")
+
+        if not test_models():
+            raise Exception("Model testing failed")
+
+        print("Model creation and testing completed successfully.")
+    except Exception as e:
+        print(f"Critical error in main execution: {str(e)}")

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+mlflow
+numpy
+Pillow
+scikit-image

run_colorizer.py

@@ -0,0 +1,118 @@
+import argparse
+import os
+import torch
+import mlflow
+from data_ingestion import create_dataloaders, test_data_ingestion
+from model import Generator, Discriminator, init_weights, test_models
+from train import train, test_training
+from app import setup_gradio_app
+
+EXPERIMENT_NAME = "Colorizer_Experiment"
+
+def setup_mlflow():
+    experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
+    else:
+        experiment_id = experiment.experiment_id
+    return experiment_id
+
+def run_pipeline(args):
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    experiment_id = setup_mlflow()
+
+    if args.ingest_data or args.run_all:
+        print("Starting data ingestion...")
+        train_loader = create_dataloaders(batch_size=args.batch_size)
+        if train_loader is None:
+            print("Data ingestion failed.")
+            return
+    else:
+        train_loader = None
+
+    if args.create_model or args.train or args.run_all:
+        print("Creating and testing models...")
+        generator = Generator().to(device)
+        discriminator = Discriminator().to(device)
+        generator.apply(init_weights)
+        discriminator.apply(init_weights)
+        if not test_models():
+            print("Model creation or testing failed.")
+            return
+    else:
+        generator = None
+        discriminator = None
+
+    if args.train or args.run_all:
+        print("Starting model training...")
+        if train_loader is None:
+            print("Creating dataloader for training...")
+            train_loader = create_dataloaders(batch_size=args.batch_size)
+            if train_loader is None:
+                print("Failed to create dataloader for training.")
+                return
+        if generator is None or discriminator is None:
+            print("Creating models for training...")
+            generator = Generator().to(device)
+            discriminator = Discriminator().to(device)
+            generator.apply(init_weights)
+            discriminator.apply(init_weights)
+        run_id = train(generator, discriminator, train_loader, num_epochs=args.num_epochs, device=device)
+        if run_id:
+            print(f"Training completed. Run ID: {run_id}")
+            with open("latest_run_id.txt", "w") as f:
+                f.write(run_id)
+        else:
+            print("Training failed.")
+            return
+
+    if args.test_training:
+        print("Testing training process...")
+        if train_loader is None:
+            print("Creating dataloader for testing...")
+            train_loader = create_dataloaders(batch_size=args.batch_size)
+            if train_loader is None:
+                print("Failed to create dataloader for testing.")
+                return
+        if generator is None or discriminator is None:
+            print("Creating models for testing...")
+            generator = Generator().to(device)
+            discriminator = Discriminator().to(device)
+            generator.apply(init_weights)
+            discriminator.apply(init_weights)
+        if test_training(generator, discriminator, train_loader, device):
+            print("Training process test passed.")
+        else:
+            print("Training process test failed.")
+
+    if args.serve or args.run_all:
+        print("Setting up Gradio app for serving...")
+        if not args.run_id:
+            try:
+                with open("latest_run_id.txt", "r") as f:
+                    args.run_id = f.read().strip()
+            except FileNotFoundError:
+                print("No run ID provided and couldn't find latest_run_id.txt")
+                return
+        iface = setup_gradio_app(args.run_id, device)
+        iface.launch(share=args.share)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run Colorizer Pipeline")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
+                        help="Device to use (cuda/cpu)")
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for training")
+    parser.add_argument("--num_epochs", type=int, default=50, help="Number of epochs to train")
+    parser.add_argument("--run_id", type=str, help="MLflow run ID of the trained model for inference")
+    parser.add_argument("--ingest_data", action="store_true", help="Run data ingestion")
+    parser.add_argument("--create_model", action="store_true", help="Create and test the model")
+    parser.add_argument("--train", action="store_true", help="Train the model")
+    parser.add_argument("--test_training", action="store_true", help="Test the training process")
+    parser.add_argument("--serve", action="store_true", help="Serve the model using Gradio")
+    parser.add_argument("--run_all", action="store_true", help="Run all steps")
+    parser.add_argument("--share", action="store_true", help="Share the Gradio app publicly")
+    args = parser.parse_args()
+
+    run_pipeline(args)

train.py

@@ -0,0 +1,235 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision.utils import make_grid
+import mlflow
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import numpy as np
+from skimage.color import lab2rgb, rgb2lab
+import argparse
+from itertools import islice
+from PIL import Image
+import torchvision.transforms as transforms
+
+from data_ingestion import ColorizeIterableDataset, create_dataloaders
+from model import Generator, Discriminator, init_weights
+
+EXPERIMENT_NAME = "Colorizer_Experiment"
+
+def setup_mlflow():
+    experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
+    else:
+        experiment_id = experiment.experiment_id
+    return experiment_id
+
+def lab_to_rgb(L, ab):
+    """Convert L and ab channels to RGB image"""
+    L = (L + 1.) * 50.
+    ab = ab * 128.
+    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
+    rgb_imgs = []
+    for img in Lab:
+        img_rgb = lab2rgb(img)
+        rgb_imgs.append(img_rgb)
+    return np.stack(rgb_imgs, axis=0)
+
+def preprocess_image(image_path):
+    img = Image.open(image_path).convert('RGB')
+    img = img.resize((256, 256))  # Resize to a consistent size
+    img_lab = rgb2lab(img)
+    img_lab = (img_lab + [0, 128, 128]) / [100, 255, 255]  # Normalize LAB values
+    return img_lab[:,:,0], img_lab[:,:,1:]
+
+def visualize_results(epoch, generator, train_loader, device):
+    generator.eval()
+    with torch.no_grad():
+        for inputs, real_AB in train_loader:
+            inputs, real_AB = inputs.to(device), real_AB.to(device)
+            fake_AB = generator(inputs)
+            
+            fake_rgb = lab_to_rgb(inputs.cpu(), fake_AB.cpu())
+            real_rgb = lab_to_rgb(inputs.cpu(), real_AB.cpu())
+            
+            img_grid = make_grid(torch.from_numpy(np.concatenate([real_rgb, fake_rgb], axis=3)).permute(0, 3, 1, 2), normalize=True, nrow=4)
+            
+            plt.figure(figsize=(15, 15))
+            plt.imshow(img_grid.permute(1, 2, 0).cpu())
+            plt.axis('off')
+            plt.title(f'Epoch {epoch}')
+            plt.savefig(f'results/epoch_{epoch}.png')
+            mlflow.log_artifact(f'results/epoch_{epoch}.png')
+            plt.close()
+            break  # Only visualize one batch
+    generator.train()
+
+def save_checkpoint(state, filename="checkpoint.pth.tar"):
+    torch.save(state, filename)
+    mlflow.log_artifact(filename)
+
+def load_checkpoint(filename, generator, discriminator, optimizerG, optimizerD):
+    if os.path.isfile(filename):
+        print(f"Loading checkpoint '{filename}'")
+        checkpoint = torch.load(filename)
+        start_epoch = checkpoint['epoch'] + 1
+        generator.load_state_dict(checkpoint['generator_state_dict'])
+        discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
+        optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])
+        optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])
+        print(f"Loaded checkpoint '{filename}' (epoch {checkpoint['epoch']})")
+        return start_epoch
+    else:
+        print(f"No checkpoint found at '{filename}'")
+        return 0
+
+def train(generator, discriminator, train_loader, num_epochs, device, lr=0.0002, beta1=0.5):
+    criterion = nn.BCEWithLogitsLoss()
+    l1_loss = nn.L1Loss()
+
+    optimizerG = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
+    optimizerD = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))
+
+    checkpoint_dir = "checkpoints"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    os.makedirs("results", exist_ok=True)
+
+    checkpoint_path = os.path.join(checkpoint_dir, "latest_checkpoint.pth.tar")
+    start_epoch = load_checkpoint(checkpoint_path, generator, discriminator, optimizerG, optimizerD)
+
+    experiment_id = setup_mlflow()
+    with mlflow.start_run(experiment_id=experiment_id, run_name="training_run") as run:
+        try:
+            for epoch in range(start_epoch, num_epochs):
+                generator.train()
+                discriminator.train()
+
+                # Use a fixed number of iterations per epoch
+                num_iterations = 1000
+                pbar = tqdm(enumerate(islice(train_loader, num_iterations)), total=num_iterations, desc=f"Epoch {epoch+1}/{num_epochs}")
+                
+                for i, (real_L, real_AB) in pbar:
+                    real_L, real_AB = real_L.to(device), real_AB.to(device)
+                    batch_size = real_L.size(0)
+
+                    # Train Discriminator
+                    optimizerD.zero_grad()
+
+                    fake_AB = generator(real_L)
+                    fake_LAB = torch.cat([real_L, fake_AB], dim=1)
+                    real_LAB = torch.cat([real_L, real_AB], dim=1)
+
+                    pred_fake = discriminator(fake_LAB.detach())
+                    loss_D_fake = criterion(pred_fake, torch.zeros_like(pred_fake))
+
+                    pred_real = discriminator(real_LAB)
+                    loss_D_real = criterion(pred_real, torch.ones_like(pred_real))
+
+                    loss_D = (loss_D_fake + loss_D_real) * 0.5
+                    loss_D.backward()
+                    optimizerD.step()
+
+                    # Train Generator
+                    optimizerG.zero_grad()
+
+                    fake_AB = generator(real_L)
+                    fake_LAB = torch.cat([real_L, fake_AB], dim=1)
+                    pred_fake = discriminator(fake_LAB)
+
+                    loss_G_GAN = criterion(pred_fake, torch.ones_like(pred_fake))
+                    loss_G_L1 = l1_loss(fake_AB, real_AB) * 100  # L1 loss weight
+
+                    loss_G = loss_G_GAN + loss_G_L1
+                    loss_G.backward()
+                    optimizerG.step()
+
+                    pbar.set_postfix({
+                        'D_loss': loss_D.item(),
+                        'G_loss': loss_G.item(),
+                        'G_L1': loss_G_L1.item()
+                    })
+
+                    mlflow.log_metrics({
+                        "D_loss": loss_D.item(),
+                        "G_loss": loss_G.item(),
+                        "G_L1_loss": loss_G_L1.item()
+                    }, step=epoch * num_iterations + i)
+
+                visualize_results(epoch, generator, train_loader, device)
+
+                checkpoint = {
+                    'epoch': epoch,
+                    'generator_state_dict': generator.state_dict(),
+                    'discriminator_state_dict': discriminator.state_dict(),
+                    'optimizerG_state_dict': optimizerG.state_dict(),
+                    'optimizerD_state_dict': optimizerD.state_dict(),
+                }
+                save_checkpoint(checkpoint, filename=checkpoint_path)
+
+            print("Training completed successfully.")
+            
+            # Log the generator model
+            mlflow.pytorch.log_model(generator, "generator_model")
+            
+            # Register the model
+            model_uri = f"runs:/{run.info.run_id}/generator_model"
+            mlflow.register_model(model_uri, "colorizer_generator")
+            
+            return run.info.run_id
+            
+        except Exception as e:
+            print(f"Error during training: {str(e)}")
+            mlflow.log_param("error", str(e))
+            return None
+
+def test_training(generator, discriminator, train_loader, device):
+    print("Testing training process...")
+    try:
+        train(generator, discriminator, train_loader, num_epochs=1, device=device)
+        print("Training process test passed.")
+        return True
+    except Exception as e:
+        print(f"Training process test failed: {str(e)}")
+        return False
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train Colorizer model")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
+                        help="Device to use for training (cuda/cpu)")
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for training")
+    parser.add_argument("--num_epochs", type=int, default=50, help="Number of epochs to train")
+    parser.add_argument("--test", action="store_true", help="Run in test mode")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    try:
+        train_loader = create_dataloaders(batch_size=args.batch_size)
+
+        generator = Generator().to(device)
+        discriminator = Discriminator().to(device)
+
+        generator.apply(init_weights)
+        discriminator.apply(init_weights)
+
+        if args.test:
+            if test_training(generator, discriminator, train_loader, device):
+                print("All tests passed.")
+            else:
+                print("Tests failed.")
+        else:
+            run_id = train(generator, discriminator, train_loader, num_epochs=args.num_epochs, device=device)
+            if run_id:
+                print(f"Training completed. Run ID: {run_id}")
+                # Save the run ID to a file for easy access by the inference script
+                with open("latest_run_id.txt", "w") as f:
+                    f.write(run_id)
+            else:
+                print("Training failed.")
+
+    except Exception as e:
+        print(f"Critical error in main execution: {str(e)}")