Add cv2 to the requirements

pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
     "numpy>=2.2.4",
+    "opencv-python>=4.11.0.86",
     "pandas>=2.2.3",
     "torch>=2.6.0",
     "torchvision>=0.21.0",

requirements.txt

@@ -26,6 +26,7 @@ nvidia-cusparselt-cu12==0.6.2
 nvidia-nccl-cu12==2.21.5
 nvidia-nvjitlink-cu12==12.4.127
 nvidia-nvtx-cu12==12.4.127
+opencv-python==4.11.0.86
 pandas==2.2.3
 pillow==11.1.0
 platformdirs==4.3.7

training_scripts/train_ddpm.py

@@ -1,305 +0,0 @@
-import logging
-import os
-
-import numpy as np
-import torch
-import wandb
-import yaml
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from data import mnist_dataset
-from data.mnist_dataset import MnistDataset
-from model_blocks.unet_base import UNet
-from scheduler.linear_scheduler import LinearNoiseScheduler
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-logger = logging.getLogger(__name__)
-wandb.login()
-
-
-def init_wandb(config):
-    """
-    Initialize a new wandb run
-    """
-    run = wandb.init(
-        project="controlnet-ddpm-mnist",
-        config=config,
-        resume="allow",  # Allows resuming if run was interrupted
-    )
-    return run
-
-
-def load_checkpoint(model, optimizer, scheduler, checkpoint_path):
-    """
-    Load model checkpoint from local file
-    """
-    checkpoint = torch.load(checkpoint_path, map_location=device)
-
-    model.load_state_dict(checkpoint["model_state_dict"])
-    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
-    if (
-        scheduler
-        and "scheduler_state_dict" in checkpoint
-        and checkpoint["scheduler_state_dict"]
-    ):
-        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
-
-    start_epoch = checkpoint["epoch"] + 1  # Start from the next epoch
-    step = checkpoint["step"]
-
-    print(f"Loaded checkpoint from epoch {checkpoint['epoch']}")
-    return start_epoch, step
-
-
-def save_checkpoint(
-    model, optimizer, scheduler, epoch, loss, step, run, checkpoint_path
-):
-    """
-    Save model checkpoint locally and to wandb
-    """
-    # Create checkpoint dictionary
-    checkpoint = {
-        "epoch": epoch,
-        "model_state_dict": model.state_dict(),
-        "optimizer_state_dict": optimizer.state_dict(),
-        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
-        "loss": loss,
-        "step": step,
-    }
-
-    # Save locally
-    torch.save(checkpoint, checkpoint_path)
-
-    # Log to wandb
-    artifact = wandb.Artifact(f"model-checkpoint-epoch-{epoch}", type="model")
-    artifact.add_file(checkpoint_path)
-    run.log_artifact(artifact)
-
-    print(f"Checkpoint saved at epoch {epoch}")
-    return checkpoint_path
-
-
-def train(args):
-    with open(args.config_path, "r") as file:
-        try:
-            config = yaml.safe_load(file)
-        except yaml.YAMLError as exc:
-            print(exc)
-    print(config)
-    run = init_wandb(config["train_params"])
-
-    diffusion_config = config["diffusion_params"]
-    dataset_config = config["dataset_params"]
-    model_config = config["model_params"]
-    train_config = config["train_params"]
-    scheduler = LinearNoiseScheduler(
-        num_timesteps=diffusion_config["num_timesteps"],
-        beta_start=diffusion_config["beta_start"],
-        beta_end=diffusion_config["beta_end"],
-    )
-
-    mnist = MnistDataset("train", dataset_config["im_path"])
-    mnist_loader = DataLoader(
-        mnist, batch_size=train_config["batch_size"], shuffle=True, num_workers=4
-    )
-
-    model = UNet(model_config).to(device)
-    model.train()
-    logger.debug(f"Initialized model and set to train")
-    optimizer = torch.optim.AdamW(model.parameters(), lr=train_config["ddpm_lr"])
-    criterion = torch.nn.MSELoss()
-
-    # Create the output directories
-    if not os.path.exists(
-        os.path.join(train_config["task_name"], train_config["ddpm_ckpt_name"])
-    ):
-        os.mkdir(train_config["task_name"])
-
-    ckpt_path = os.path.join(train_config["task_name"], train_config["ddpm_ckpt_name"])
-    # Load checkpoint if there
-    if os.path.exists(ckpt_path):
-        start_epoch, step = load_checkpoint(
-            model, optimizer=optimizer, scheduler=scheduler, checkpoint_path=ckpt_path
-        )
-    else:
-        start_epoch = 0
-        step = 0
-
-    # Log model architecture as a Table
-    model_table = wandb.Table(columns=["Layer", "Parameters"])
-    total_params = 0
-    for name, param in model.named_parameters():
-        if param.requires_grad:
-            params = param.numel()
-            total_params += params
-            model_table.add_data(name, params)
-
-    wandb.log({"model_architecture": model_table})
-    wandb.log({"total_parameters": total_params})
-
-    # Watch model gradients and parameters
-    wandb.watch(model, log="all", log_freq=100)
-
-    for epoch in range(start_epoch, train_config["num_epochs"]):
-        losses = []
-        batch_idx = 0
-        progress_bar = tqdm(
-            mnist_loader, desc=f"Epoch {epoch + 1}/{train_config['num_epochs']}"
-        )
-
-        for im in progress_bar:
-            batch_idx += 1
-            optimizer.zero_grad()
-            im = im.float().to(device)
-
-            # Sample noise
-            noise = torch.randn_like(im).to(device)
-            logger.debug(f"Sampled noise epoch {epoch} : {noise.shape}")
-
-            # Sample timestep
-            t = torch.randint(0, diffusion_config["num_timesteps"], (im.shape[0],)).to(
-                device
-            )
-
-            noisy_im = scheduler.add_noise(im, noise, t)
-            noise_pred = model(noisy_im, t)
-
-            loss = criterion(noise_pred, noise)
-            losses.append(loss.item())
-            loss.backward()
-            optimizer.step()
-
-            # Calculate gradient norm for monitoring
-            total_norm = 0
-            for p in model.parameters():
-                if p.grad is not None:
-                    param_norm = p.grad.data.norm(2)
-                    total_norm += param_norm.item() ** 2
-            total_norm = total_norm**0.5
-
-            # Update progress bar
-            progress_bar.set_postfix({"loss": loss.item(), "avg_loss": np.mean(losses)})
-
-            wandb.log(
-                {
-                    "train/batch_loss": loss.item(),
-                    "train/step": step,
-                    "train/epoch": epoch + batch_idx / len(mnist_loader),
-                    "train/gradient_norm": total_norm,
-                    "train/learning_rate": optimizer.param_groups[0]["lr"],
-                }
-            )
-
-            step += 1
-
-        avg_loss = np.mean(losses)
-
-        # Log epoch-level metrics
-        wandb.log(
-            {
-                "train/epoch_loss": avg_loss,
-                "train/epoch_completed": epoch,
-            }
-        )
-
-        print(f"Finished epoch: {epoch} | Loss: {np.mean(losses):.4f}")
-        if epoch % train_config["save_epoch"]:
-            visualize_samples(
-                model, scheduler, epoch, diffusion_config["num_timesteps"], device
-            )
-
-        save_checkpoint(
-            model, optimizer, scheduler, epoch, np.mean(losses), step, run, ckpt_path
-        )
-
-    # Log final model as artifact
-    logging.info("Finished training and starting to save model")
-    final_model_path = os.path.join(
-        train_config["task_name"], f"final_{train_config['ddpm_ckpt_name']}"
-    )
-    save_checkpoint(
-        model,
-        optimizer,
-        scheduler,
-        train_config["num_epochs"] - 1,
-        avg_loss,
-        step,
-        run,
-        final_model_path,
-    )
-    logging.info("Saved Model to Wandb and local")
-
-    # Log a summary table of training
-    summary_table = wandb.Table(columns=["Metric", "Value"])
-    summary_table.add_data("Final Loss", avg_loss)
-    summary_table.add_data("Best Loss", best_loss)
-    summary_table.add_data("Best Epoch", epoch)
-    summary_table.add_data("Total Steps", step)
-    summary_table.add_data("Training Time (hours)", wandb.run.duration / 3600)
-
-    wandb.log({"training_summary": summary_table})
-
-    # Finish the run
-    wandb.finish()
-
-
-def visualize_samples(model, scheduler, epoch, num_timesteps, device, num_samples=4):
-    """
-    Generate sample images from noise and log to wandb
-    """
-    model.eval()
-    with torch.no_grad():
-        # Start with random noise
-        samples = torch.randn(num_samples, 1, 28, 28).to(device)
-
-        # Store the denoising process
-        sample_images = []
-
-        # Record more frequently at the beginning of sampling
-        log_steps = set([0, 20, 50, 100, 200, 400, 600, 800, num_timesteps - 1])
-
-        # Denoise gradually
-        for i in tqdm(reversed(range(num_timesteps)), desc="Sampling"):
-            t = torch.full((num_samples,), i, device=device, dtype=torch.long)
-
-            # Get model prediction and update sample
-            predicted_noise = model(samples, t)
-            samples = scheduler.step(predicted_noise, i, samples)
-
-            # Save images at specified timesteps
-            if i in log_steps:
-                # Denormalize and convert to numpy for logging
-                denorm_samples = samples.clamp(-1, 1).cpu().numpy()
-                denorm_samples = (
-                    denorm_samples + 1
-                ) / 2.0  # scale from [-1, 1] to [0, 1]
-                sample_images.append((i, denorm_samples))
-
-    # Create a grid to show denoising process
-    images_to_log = {}
-
-    # Log individual samples
-    for i, sample in enumerate(samples):
-        sample_np = sample.clamp(-1, 1).cpu().numpy()
-        sample_np = (sample_np + 1) / 2.0  # scale from [-1, 1] to [0, 1]
-        images_to_log[f"sample_{i}_epoch_{epoch}"] = wandb.Image(
-            sample_np[0], caption=f"Sample {i}, Epoch {epoch}"
-        )
-
-    # Log denoising process for first sample
-    denoising_steps = []
-    for step_idx, samples_np in sample_images:
-        denoising_steps.append(
-            wandb.Image(
-                samples_np[0][0],
-                caption=f"Step {num_timesteps - step_idx}/{num_timesteps}",
-            )
-        )
-
-    images_to_log["denoising_process_epoch_" + str(epoch)] = denoising_steps
-
-    # Log all images
-    wandb.log(images_to_log)
-
-    model.train()

uv.lock

@@ -2,8 +2,10 @@ version = 1
 revision = 1
 requires-python = ">=3.13"
 resolution-markers = [
-    "sys_platform == 'linux'",
-    "sys_platform != 'linux'",
+    "platform_machine == 'aarch64' and sys_platform == 'linux'",
+    "platform_machine != 'aarch64' and sys_platform == 'linux'",
+    "sys_platform == 'darwin'",
+    "sys_platform != 'darwin' and sys_platform != 'linux'",
 ]
 
 [[package]]
@@ -73,6 +75,7 @@ version = "0.1.0"
 source = { virtual = "." }
 dependencies = [
     { name = "numpy" },
+    { name = "opencv-python" },
     { name = "pandas" },
     { name = "torch" },
     { name = "torchvision" },
@@ -83,6 +86,7 @@ dependencies = [
 [package.metadata]
 requires-dist = [
     { name = "numpy", specifier = ">=2.2.4" },
+    { name = "opencv-python", specifier = ">=4.11.0.86" },
     { name = "pandas", specifier = ">=2.2.3" },
     { name = "torch", specifier = ">=2.6.0" },
     { name = "torchvision", specifier = ">=0.21.0" },
@@ -276,7 +280,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.1.0.70"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
@@ -287,7 +291,7 @@ name = "nvidia-cufft-cu12"
 version = "11.2.1.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117 },
@@ -306,9 +310,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.6.1.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
-    { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" },
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
+    { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057 },
@@ -319,7 +323,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.3.1.170"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763 },
@@ -357,6 +361,23 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144 },
 ]
 
+[[package]]
+name = "opencv-python"
+version = "4.11.0.86"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/17/06/68c27a523103dad5837dc5b87e71285280c4f098c60e4fe8a8db6486ab09/opencv-python-4.11.0.86.tar.gz", hash = "sha256:03d60ccae62304860d232272e4a4fda93c39d595780cb40b161b310244b736a4", size = 95171956 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/05/4d/53b30a2a3ac1f75f65a59eb29cf2ee7207ce64867db47036ad61743d5a23/opencv_python-4.11.0.86-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:432f67c223f1dc2824f5e73cdfcd9db0efc8710647d4e813012195dc9122a52a", size = 37326322 },
+    { url = "https://files.pythonhosted.org/packages/3b/84/0a67490741867eacdfa37bc18df96e08a9d579583b419010d7f3da8ff503/opencv_python-4.11.0.86-cp37-abi3-macosx_13_0_x86_64.whl", hash = "sha256:9d05ef13d23fe97f575153558653e2d6e87103995d54e6a35db3f282fe1f9c66", size = 56723197 },
+    { url = "https://files.pythonhosted.org/packages/f3/bd/29c126788da65c1fb2b5fb621b7fed0ed5f9122aa22a0868c5e2c15c6d23/opencv_python-4.11.0.86-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b92ae2c8852208817e6776ba1ea0d6b1e0a1b5431e971a2a0ddd2a8cc398202", size = 42230439 },
+    { url = "https://files.pythonhosted.org/packages/2c/8b/90eb44a40476fa0e71e05a0283947cfd74a5d36121a11d926ad6f3193cc4/opencv_python-4.11.0.86-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b02611523803495003bd87362db3e1d2a0454a6a63025dc6658a9830570aa0d", size = 62986597 },
+    { url = "https://files.pythonhosted.org/packages/fb/d7/1d5941a9dde095468b288d989ff6539dd69cd429dbf1b9e839013d21b6f0/opencv_python-4.11.0.86-cp37-abi3-win32.whl", hash = "sha256:810549cb2a4aedaa84ad9a1c92fbfdfc14090e2749cedf2c1589ad8359aa169b", size = 29384337 },
+    { url = "https://files.pythonhosted.org/packages/a4/7d/f1c30a92854540bf789e9cd5dde7ef49bbe63f855b85a2e6b3db8135c591/opencv_python-4.11.0.86-cp37-abi3-win_amd64.whl", hash = "sha256:085ad9b77c18853ea66283e98affefe2de8cc4c1f43eda4c100cf9b2721142ec", size = 39488044 },
+]
+
 [[package]]
 name = "pandas"
 version = "2.2.3"