""" Based on https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py """ from typing import Any, Dict import numpy as np from .gaussian_diffusion import ( GaussianDiffusion, SpacedDiffusion, get_named_beta_schedule, space_timesteps, ) BASE_DIFFUSION_CONFIG = { "channel_biases": [0.0, 0.0, 0.0, -1.0, -1.0, -1.0], "channel_scales": [2.0, 2.0, 2.0, 0.007843137255, 0.007843137255, 0.007843137255], "mean_type": "epsilon", "schedule": "cosine", "timesteps": 1024, } DIFFUSION_CONFIGS = { "base40M-imagevec": BASE_DIFFUSION_CONFIG, "base40M-textvec": BASE_DIFFUSION_CONFIG, "base40M-uncond": BASE_DIFFUSION_CONFIG, "base40M": BASE_DIFFUSION_CONFIG, "base300M": BASE_DIFFUSION_CONFIG, "base1B": BASE_DIFFUSION_CONFIG, "upsample": { "channel_biases": [0.0, 0.0, 0.0, -1.0, -1.0, -1.0], "channel_scales": [2.0, 2.0, 2.0, 0.007843137255, 0.007843137255, 0.007843137255], "mean_type": "epsilon", "schedule": "linear", "timesteps": 1024, }, } def diffusion_from_config(config: Dict[str, Any]) -> GaussianDiffusion: schedule = config["schedule"] steps = config["timesteps"] respace = config.get("respacing", None) mean_type = config.get("mean_type", "epsilon") betas = get_named_beta_schedule(schedule, steps) channel_scales = config.get("channel_scales", None) channel_biases = config.get("channel_biases", None) if channel_scales is not None: channel_scales = np.array(channel_scales) if channel_biases is not None: channel_biases = np.array(channel_biases) kwargs = dict( betas=betas, model_mean_type=mean_type, model_var_type="learned_range", loss_type="mse", channel_scales=channel_scales, channel_biases=channel_biases, ) if respace is None: return GaussianDiffusion(**kwargs) else: return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace), **kwargs)