mhc-ddim 0.1.0__tar.gz

mhc_ddim-0.1.0/LICENSE

@@ -0,0 +1 @@
+Copyright (c) 2026 Yasin Ansari

mhc_ddim-0.1.0/PKG-INFO

@@ -0,0 +1,27 @@
+Metadata-Version: 2.4
+Name: mhc-ddim
+Version: 0.1.0
+Summary: DDIM and DDPM implementations with Manifold-Constrained Hyper-Connections (mHC)
+Author-email: Yasin Ansari <yasinansari.work@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/TryingManV2/mhc-ddim
+Project-URL: Repository, https://github.com/TryingManV2/mhc-ddim
+Project-URL: Issues, https://github.com/TryingManV2/mhc-ddim/issues
+Keywords: diffusion,ddpm,ddim,mHC,manifold-constrained hyper-connections,pytorch,mnist,deep-learning,generative-models
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: matplotlib
+Requires-Dist: tqdm
+Dynamic: license-file

mhc_ddim-0.1.0/pyproject.toml

@@ -0,0 +1,55 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mhc-ddim"
+version = "0.1.0"
+description = "DDIM and DDPM implementations with Manifold-Constrained Hyper-Connections (mHC)"
+readme = "README.md"
+requires-python = ">=3.9"
+
+authors = [
+    { name = "Yasin Ansari", email = "yasinansari.work@gmail.com" }
+]
+
+license = { text = "MIT" }
+
+keywords = [
+    "diffusion",
+    "ddpm",
+    "ddim",
+    "mHC",
+    "manifold-constrained hyper-connections",
+    "pytorch",
+    "mnist",
+    "deep-learning",
+    "generative-models"
+]
+
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+    "torch",
+    "numpy",
+    "matplotlib",
+    "tqdm",
+]
+
+[project.urls]
+Homepage = "https://github.com/TryingManV2/mhc-ddim"
+Repository = "https://github.com/TryingManV2/mhc-ddim"
+Issues = "https://github.com/TryingManV2/mhc-ddim/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]

mhc_ddim-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mhc_ddim-0.1.0/src/mhc_ddim/cfg.py

@@ -0,0 +1,48 @@
+import torch
+
+
+class CFGModelWrapper:
+    """
+    Wraps a conditional UNet to implement classifier-free guidance at sampling.
+
+    Runs the model twice per step — once with the given class labels, once
+    with the null class — then combines:
+
+        eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
+
+    Set guidance_scale = 1.0 to recover plain conditional sampling.
+    Set guidance_scale = 0.0 to recover plain unconditional sampling.
+
+    Only valid for ModelMeanType.EPSILON. For START_X or PREVIOUS_X you
+    would need the analogous combination on those quantities.
+    """
+    def __init__(self, model, null_class, guidance_scale=3.0):
+        self.model = model
+        self.null_class = null_class
+        self.guidance_scale = guidance_scale
+
+    def __call__(self, x, t, **kwargs):
+        y = kwargs.get("y", None)
+        if y is None:
+            # No label provided: behave like the plain unconditional model.
+            return self.model(x, t, **kwargs)
+
+        null_y = torch.full_like(y, self.null_class)
+
+        kwargs_cond   = {**kwargs, "y": y}
+        kwargs_uncond = {**kwargs, "y": null_y}
+
+        out_cond   = self.model(x, t, **kwargs_cond)
+        out_uncond = self.model(x, t, **kwargs_uncond)
+
+        # If the model outputs variance too (learn_sigma=True), the second
+        # half of the channels must NOT be guided. Split and re-combine.
+        if out_cond.shape[1] == 2 * x.shape[1]:
+            eps_c, var_c = out_cond.chunk(2, dim=1)
+            eps_u, var_u = out_uncond.chunk(2, dim=1)
+            eps = eps_u + self.guidance_scale * (eps_c - eps_u)
+            # Use conditional variance; using unconditional is also fine.
+            return torch.cat([eps, var_c], dim=1)
+
+        # Standard case: pure epsilon prediction.
+        return out_uncond + self.guidance_scale * (out_cond - out_uncond)

mhc_ddim-0.1.0/src/mhc_ddim/dist_util.py

@@ -0,0 +1,85 @@
+import io
+import os
+import socket
+
+import blobfile as bf # type: ignore
+from mpi4py import MPI # type: ignore
+import torch 
+import torch.distributed as dist
+
+# Change this to reflect your cluster layout.
+# The GPU for a given rank is (rank % GPUS_PER_NODE).
+GPUS_PER_NODE = 8
+
+SETUP_RETRY_COUNT = 3
+
+
+def setup_dist():
+    """
+    Setup a distributed process group.
+    """
+    if dist.is_initialized():
+        return
+
+    comm = MPI.COMM_WORLD
+    backend = "gloo" if not torch.cuda.is_available() else "nccl"
+
+    if backend == "gloo":
+        hostname = "localhost"
+    else:
+        hostname = socket.gethostbyname(socket.getfqdn())
+    os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
+    os.environ["RANK"] = str(comm.rank)
+    os.environ["WORLD_SIZE"] = str(comm.size)
+
+    port = comm.bcast(_find_free_port(), root=0)
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group(backend=backend, init_method="env://")
+
+
+def dev():
+    """
+    Get the device to use for torch.distributed.
+    """
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}")
+    return torch.device("cpu")
+
+
+def load_state_dict(path, **kwargs):
+    """
+    Load a PyTorch file without redundant fetches across MPI ranks.
+    """
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        with bf.BlobFile(path, "rb") as f:
+            data = f.read()
+    else:
+        data = None
+    data = MPI.COMM_WORLD.bcast(data)
+
+    # PyTorch >= 2.6 defaults `weights_only` to True, which rejects the
+    # Python objects stored in optimizer checkpoints. We load our own
+    # checkpoints, so explicitly fall back to the permissive unpickler
+    # unless the caller already chose otherwise.
+    kwargs.setdefault("weights_only", False)
+
+    return torch.load(io.BytesIO(data), **kwargs)
+
+
+def sync_params(params):
+    """
+    Synchronize a sequence of Tensors across ranks from rank 0.
+    """
+    for p in params:
+        with torch.no_grad():
+            dist.broadcast(p, 0)
+
+
+def _find_free_port():
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+    finally:
+        s.close()

mhc_ddim-0.1.0/src/mhc_ddim/fp16_util.py

@@ -0,0 +1,69 @@
+import torch.nn as nn
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+
+def convert_module_to_f16(l):
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        assert l.weight is not None
+        l.weight.data = l.weight.data.half()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.half()
+
+def convert_module_to_f32(l):
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        assert l.weight is not None
+        l.weight.data = l.weight.data.float()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.float()
+
+
+def make_master_params(model_params):
+    """
+    Copy model parameters into a (differently-shaped) list of full-precision
+    parameters.
+    """
+    master_params = _flatten_dense_tensors(
+        [param.detach().float() for param in model_params]
+    )
+    master_params = nn.Parameter(master_params)
+    master_params.requires_grad = True
+    return [master_params]
+
+
+def model_grads_to_master_grads(model_params, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    master_params[0].grad = _flatten_dense_tensors(
+        [param.grad.data.detach().float() for param in model_params]
+    )
+
+
+def master_params_to_model_params(model_params, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    # Without copying to a list, if a generator is passed, this will
+    # silently not copy any parameters.
+    model_params = list(model_params)
+
+    for param, master_param in zip(
+        model_params, unflatten_master_params(model_params, master_params)
+    ):
+        param.detach().copy_(master_param)
+
+
+def unflatten_master_params(model_params, master_params):
+    """
+    Unflatten the master parameters to look like model_params.
+    """
+    return _unflatten_dense_tensors(master_params[0].detach(), tuple(tensor for tensor in model_params))
+
+
+def zero_grad(model_params):
+    for param in model_params:
+        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
+        if param.grad is not None:
+            param.grad.detach_()
+            param.grad.zero_()