diffusionmodel.py

# -*- coding: utf-8 -*-
"""DiffusionModel.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1PheRVtX0V3mS1BklbLKhqSIZgUuYQ4Af
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os
from tqdm import tqdm

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=3, time_dim=256, base_channels=64):
        super().__init__()
        self.time_dim = time_dim

        # Time embedding
        self.time_mlp = nn.Sequential(
            nn.Linear(1, time_dim),
            nn.SiLU(),
            nn.Linear(time_dim, time_dim)
        )

        # Encoder (downsampling)
        self.inc = self._make_input_block(in_channels, base_channels)
        self.down1 = self._make_down_block(base_channels, base_channels*2)
        self.down2 = self._make_down_block(base_channels*2, base_channels*4)
        self.down3 = self._make_down_block(base_channels*4, base_channels*8)

        # Bottleneck
        self.bottleneck = nn.Sequential(
            nn.Conv2d(base_channels*8, base_channels*8, 3, padding=1),
            nn.BatchNorm2d(base_channels*8),
            nn.ReLU(),
            nn.Conv2d(base_channels*8, base_channels*8, 3, padding=1),
            nn.BatchNorm2d(base_channels*8),
            nn.ReLU()
        )

        # Timestep embedding projection
        self.time_proj = nn.Sequential(
            nn.Linear(time_dim, base_channels*8),
            nn.SiLU()
        )

        # Decoder (upsampling)
        self.up1 = self._make_up_block(base_channels*8, base_channels*4)
        self.up2 = self._make_up_block(base_channels*4, base_channels*2)
        self.up3 = self._make_up_block(base_channels*2, base_channels)

        # Output layer
        self.outc = nn.Sequential(
            nn.Conv2d(base_channels, out_channels, kernel_size=3, padding=1),
            nn.Tanh()  # Output in range [-1, 1]
        )

    def _make_input_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def _make_down_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def _make_up_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
            nn.Conv2d(out_channels*2, out_channels, kernel_size=3, padding=1),  # Fixed to handle skip connections
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def forward(self, x, t, condition):
        # Time embedding
        t_emb = self.time_mlp(t.unsqueeze(-1).float())

        # Concatenate input image with condition image
        x_combined = torch.cat([x, condition], dim=1)

        # Encoder
        x1 = self.inc(x_combined)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)

        # Bottleneck with timestep injection
        x4 = self.bottleneck(x4)

        # Add time embedding to bottleneck features
        t_emb = self.time_proj(t_emb)
        t_emb = t_emb.view(-1, t_emb.shape[1], 1, 1).expand(-1, -1, x4.shape[2], x4.shape[3])
        x4 = x4 + t_emb

        # Decoder with skip connections
        x = self.up1[0](x4)  # ConvTranspose only
        x = torch.cat([x, x3], dim=1)  # Skip connection
        x = self.up1[1:](x)  # Rest of the up block

        x = self.up2[0](x)
        x = torch.cat([x, x2], dim=1)
        x = self.up2[1:](x)

        x = self.up3[0](x)
        x = torch.cat([x, x1], dim=1)
        x = self.up3[1:](x)

        # Output
        return self.outc(x)

class DiffusionModel:
    def __init__(self, img_size=128, device="cuda" if torch.cuda.is_available() else "cpu"):
        self.img_size = img_size
        self.device = device
        self.model = UNet(in_channels=6, out_channels=3, base_channels=64).to(device)
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-4)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=100, gamma=0.5)

        # Define beta schedule for diffusion process
        self.timesteps = 1000
        self.beta_start = 1e-4
        self.beta_end = 0.02
        # Use cosine beta schedule instead of linear for better results
        self.betas = self._cosine_beta_schedule(self.timesteps).to(device)
        self.alphas = 1. - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
        self.sqrt_recip_alphas = torch.sqrt(1.0 / self.alphas)
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)
        self.posterior_variance = self.betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)

        # Create output directory for intermediate results
        self.output_dir = "diffusion_outputs"
        os.makedirs(self.output_dir, exist_ok=True)

    def _cosine_beta_schedule(self, timesteps, s=0.008):
        """
        Create a beta schedule that follows a cosine function.
        This helps get better results at the final steps of diffusion.
        """
        steps = timesteps + 1
        x = torch.linspace(0, timesteps, steps)
        alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        return torch.clip(betas, 0.0001, 0.9999)

    def get_transforms(self):
        return transforms.Compose([
            transforms.Resize((self.img_size, self.img_size)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

    def extract(self, a, t, x_shape):
        batch_size = t.shape[0]
        out = a.gather(-1, t.cpu()).float()
        return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

    def q_sample(self, x_start, t, noise=None):
        if noise is None:
            noise = torch.randn_like(x_start)

        sqrt_alphas_cumprod_t = self.extract(self.sqrt_alphas_cumprod, t, x_start.shape)
        sqrt_one_minus_alphas_cumprod_t = self.extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)

        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise, noise

    def train_step(self, clean_images, condition_images):
        batch_size = clean_images.shape[0]

        # Sample random timesteps
        t = torch.randint(0, self.timesteps, (batch_size,), device=self.device).long()

        # Add noise to images
        noisy_images, noise = self.q_sample(clean_images, t)

        # Predict noise using model
        predicted_noise = self.model(noisy_images, t / self.timesteps, condition_images)

        # Calculate loss (mix of L1 and L2 loss for better convergence)
        loss = F.mse_loss(predicted_noise, noise) + 0.1 * F.l1_loss(predicted_noise, noise)

        # Optimize
        self.optimizer.zero_grad()
        loss.backward()
        # Gradient clipping to prevent exploding gradients
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        self.optimizer.step()

        return loss.item()

    def generate(self, condition_image, steps=100, save_intermediate=True):
        """
        Generate an image from a condition image using the trained diffusion model
        """
        condition_image = condition_image.to(self.device)

        # Start with random noise
        img = torch.randn((1, 3, self.img_size, self.img_size), device=self.device)

        # Save initial noise
        if save_intermediate:
            self._save_image(img, "initial_noise")

        # Gradually denoise the image
        time_steps = list(reversed(range(0, self.timesteps, self.timesteps // steps)))

        with torch.no_grad():
            for i, step in enumerate(tqdm(time_steps)):
                t = torch.full((1,), step, device=self.device, dtype=torch.long)

                # Predict noise
                predicted_noise = self.model(img, t / self.timesteps, condition_image)

                # Get alpha values for this timestep
                alpha = self.alphas[step]
                alpha_cumprod = self.alphas_cumprod[step]
                beta = self.betas[step]

                if step > 0:
                    noise = torch.randn_like(img) * torch.sqrt(beta)
                else:
                    noise = torch.zeros_like(img)

                # Improved sampling formula for better results
                img = (1 / torch.sqrt(alpha)) * (
                    img - ((1 - alpha) / torch.sqrt(1 - alpha_cumprod)) * predicted_noise
                ) + noise

                # Save intermediate results at regular intervals
                if save_intermediate and i % (len(time_steps) // 10) == 0:
                    self._save_image(img, f"step_{step}")

        # Rescale image from [-1, 1] to [0, 1]
        img = (img.clamp(-1, 1) + 1) / 2

        # Save final result
        if save_intermediate:
            self._save_image(img, "final_result")

        return img

    def _save_image(self, tensor, name):
        """Helper function to save image tensors during generation"""
        img = (tensor.clamp(-1, 1) + 1) / 2
        img_np = img.cpu().squeeze().permute(1, 2, 0).numpy()
        plt.figure(figsize=(10, 10))
        plt.imshow(img_np)
        plt.axis('off')
        plt.savefig(f"{self.output_dir}/{name}.png", bbox_inches='tight', pad_inches=0)
        plt.close()

# Custom dataset for image-to-image translation
class ImagePairDataset(Dataset):
    def __init__(self, source_images, target_images, transform=None):
        self.source_images = source_images
        self.target_images = target_images
        self.transform = transform

    def __len__(self):
        return len(self.source_images)

    def __getitem__(self, idx):
        source_img = Image.open(self.source_images[idx])
        target_img = Image.open(self.target_images[idx])

        if self.transform:
            source_img = self.transform(source_img)
            target_img = self.transform(target_img)

        return target_img, source_img  # Return (clean_image, condition_image)

# Training function with progress tracking
def train_diffusion_model(dataloader, epochs=20, save_every=1):
    diffusion = DiffusionModel(img_size=128)  # Increased resolution
    transform = diffusion.get_transforms()

    # Create directories for saving models and samples
    os.makedirs("model_checkpoints", exist_ok=True)
    os.makedirs("training_samples", exist_ok=True)

    # For tracking progress
    all_losses = []

    # Get a sample condition image for generating samples during training
    sample_condition = next(iter(dataloader))[1][0:1]  # Take first condition image

    for epoch in range(epochs):
        epoch_losses = []
        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")

        for batch_idx, (clean_images, condition_images) in enumerate(progress_bar):
            clean_images = clean_images.to(diffusion.device)
            condition_images = condition_images.to(diffusion.device)

            loss = diffusion.train_step(clean_images, condition_images)
            epoch_losses.append(loss)

            # Update progress bar
            progress_bar.set_postfix(loss=f"{loss:.4f}")

        # Update learning rate
        diffusion.scheduler.step()

        # Calculate average loss for this epoch
        avg_loss = sum(epoch_losses) / len(epoch_losses)
        all_losses.append(avg_loss)
        print(f"Epoch {epoch+1} completed, Average Loss: {avg_loss:.4f}")

        # Save model checkpoint
        if (epoch + 1) % save_every == 0:
            torch.save({
                'model_state_dict': diffusion.model.state_dict(),
                'optimizer_state_dict': diffusion.optimizer.state_dict(),
                'scheduler_state_dict': diffusion.scheduler.state_dict(),
                'epoch': epoch,
                'loss': avg_loss
            }, f"model_checkpoints/diffusion_model_epoch_{epoch+1}.pth")

            # Generate and save a sample image to monitor progress
            with torch.no_grad():
                sample_image = diffusion.generate(
                    sample_condition.to(diffusion.device),
                    steps=20,  # Use fewer steps during training to save time
                    save_intermediate=False
                )

                # Save the sample
                sample_np = sample_image.cpu().squeeze().permute(1, 2, 0).numpy()
                plt.figure(figsize=(10, 10))
                plt.imshow(sample_np)
                plt.axis('off')
                plt.savefig(f"training_samples/sample_epoch_{epoch+1}.png", bbox_inches='tight')
                plt.close()

    # Plot loss curve
    plt.figure(figsize=(10, 5))
    plt.plot(all_losses)
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig('training_loss.png')
    plt.close()

    return diffusion

def example_usage():
    """
    Uses data augmentation to create variations from limited image pairs
    """
    # Original pairs (8 examples)
    source_images = [
        "/content/sunrise1.jpg",
        "/content/sunrise2.jpg",
        "/content/pexels-abdghat-1631678.jpg",
        "/content/pexels-no-name-14543-66997.jpg",
        "/content/pexels-pixabay-355508.jpg",
        "/content/pexels-sebastian-189349.jpg",
        "/content/pexels-simon73-1266810.jpg",
        "/content/pexels-simon73-1323550.jpg"
    ]

    target_images = [
        "/content/sunset1.jpg",
        "/content/sunset2.jpg",
        "/content/sunset3.jpg",
        "/content/sunset4.jpg",
        "/content/sunset5.jpg",
        "/content/sunset6.jpg",
        "/content/sunset7.jpg",
        "/content/sunset8.jpg"
    ]

    # Enhanced transformations with augmentation
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomApply([
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
            transforms.RandomRotation(10),
            transforms.RandomPerspective(distortion_scale=0.2, p=0.5)
        ], p=0.8),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # Initialize model
    diffusion = DiffusionModel(img_size=128)

    # Create augmented dataset
    class AugmentedPairDataset(Dataset):
        def __init__(self, sources, targets, transform, num_augmented=20):
            self.sources = sources
            self.targets = targets
            self.transform = transform
            self.num_augmented = num_augmented

        def __len__(self):
            return len(self.sources) * self.num_augmented

        def __getitem__(self, idx):
            pair_idx = idx % len(self.sources)
            src_img = Image.open(self.sources[pair_idx]).convert('RGB')
            tgt_img = Image.open(self.targets[pair_idx]).convert('RGB')

            # Apply same random augmentation to both images
            seed = torch.randint(0, 2**32, (1,)).item()
            torch.manual_seed(seed)
            src_img = self.transform(src_img)
            torch.manual_seed(seed)
            tgt_img = self.transform(tgt_img)

            return tgt_img, src_img  # (target, source)

    # Create dataloader (8 pairs × 20 augmentations = 160 examples per epoch)
    dataset = AugmentedPairDataset(source_images, target_images, transform, num_augmented=20)
    dataloader = DataLoader(
        dataset,
        batch_size=2,
        shuffle=True,
        num_workers=2
    )

    # Training with early stopping
    trained_model = train_diffusion_model(dataloader, epochs=50, save_every=10)

    # Generate a new image from a test condition image
    try:
        test_condition = transform(Image.open("/content/ChatGPT Image Apr 10, 2025, 04_21_16 PM.png").convert('RGB')).unsqueeze(0)
        print("Generating final image from test condition...")

        # Generate with more steps for better quality
        generated_img = trained_model.generate(
            test_condition.to(trained_model.device),
            steps=100,
            save_intermediate=True
        )

        # Convert tensor to numpy for visualization
        img_np = generated_img.cpu().squeeze().permute(1, 2, 0).numpy()

        # Save the final result
        plt.figure(figsize=(15, 15))
        plt.imshow(img_np)
        plt.axis('off')
        plt.savefig("final_generated_image.png", bbox_inches='tight', pad_inches=0)
        plt.close()

        print("Generation complete! Result saved as 'final_generated_image.png'")

    except Exception as e:
        print(f"Error during generation: {str(e)}")
        print("You can still use the trained model with a different test image.")

# Run the example
if __name__ == "__main__":
    example_usage()