Noise Schedules for Diffusion Models in C

gistfile1.txt

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <assert.h>

#define PI 3.14159265358979323846  
#define MIN_BETA 0.0001
#define MAX_BETA 0.9999

float DiffusionInternalSigmoid(float x)
{
	return 1 / (1 + exp(-x));
}

// Function to generate linearly spaced values between start and end
float *linspace(float start, float end, int arrayLength) 
{
	assert(arrayLength > 0);
	float *result = (float*)malloc(arrayLength * sizeof(float));
	assert(result != NULL);
	if(arrayLength == 1){result[0] = start;return result;}
	float step = (end - start) / (arrayLength - 1);
	for(int i = 0; i < arrayLength; i++)
	{
		result[i] = start + i * step;
	}
	return result;
}

float *LinearBetaSchedule(int timeSteps)
{
	float start = 0.0001;
	float end   = 0.02;
	return linspace(start, end, timeSteps);
}

float *QuadraticBetaSchedule(int timeSteps)
{
	float start = 0.0001;
	float end   = 0.02;
	start = sqrt(start);
	end   = sqrt(end);
	float *betas = linspace(start, end, timeSteps);
	for(int i = 0; i < timeSteps; i++)
	{
		betas[i] = betas[i] * betas[i];
	}
	return betas;
}

float *CosineBetaSchedule(int timeSteps)
{
	float s = 0.008;
	int steps = timeSteps + 1;
	float *temporary = linspace(0, timeSteps, steps);
	float *betas = calloc(timeSteps, sizeof(float));
	// Step 2: Compute alphas_cumprod using the cosine schedule formula
	float *alphasCumProd = calloc(steps, sizeof(float));
	float factor = PI * 0.5 / (1 + s);
	for(int i = 0; i < steps; i++)
	{
		float term = ((temporary[i] / timeSteps) + s) * factor;
		alphasCumProd[i] = pow(cos(term), 2);
	}
 
	// Step 3: Normalize alphas_cumprod by dividing by alphas_cumprod[0]
	float firstAlpha = alphasCumProd[0];
	for(int i = 0; i < steps; i++)
	{
		alphasCumProd[i] /= firstAlpha;
	}

	// Step 4: Compute betas
	for(int i = 1; i < steps; i++)
	{
		float beta = 1 - (alphasCumProd[i] / alphasCumProd[i - 1]);
		// Clip beta to the range [MIN_BETA, MAX_BETA]
		if(beta < MIN_BETA) beta = MIN_BETA;
		if(beta > MAX_BETA) beta = MAX_BETA;
		betas[i - 1] = beta;
	}
	free(temporary);free(alphasCumProd);
	return betas; 
}

float *SigmoidBetaSchedule(int timeSteps)
{
	float start = 0.0001;
	float end   = 0.02;
	float *betas = linspace(-6, 6, timeSteps);
	for(int i = 0; i < timeSteps; i++)
	{
		betas[i] = DiffusionInternalSigmoid(betas[i]) * (end - start) + start;
	}
	return betas;
}
void TestLinearSpace()
{
	int arrayLength = 4;
	float start = 25.7f;
	float end = 27.9f;
	float* result = linspace(start, end, arrayLength);
	assert(result != NULL);
	for(int i = 0; i < arrayLength; i++)
	{
		printf("%f ", result[i]);
	}
	free(result);
}

void TestLinearBetas()
{
	int timeSteps = 4;
	float *betas = LinearBetaSchedule(timeSteps);
	assert(betas != NULL);
	for(int i = 0; i < timeSteps; i++)
	{
		printf("%f ", betas[i]);
	}
	free(betas);
}

void TestQuadraticBetas()
{
	int timeSteps = 4;
	float *betas = QuadraticBetaSchedule(timeSteps);
	assert(betas != NULL);
	for(int i = 0; i < timeSteps; i++)
	{
		printf("%f ", betas[i]);
	}
	free(betas);
}

void TestCosineBetas()
{
	int timeSteps = 4;
	float *betas = CosineBetaSchedule(timeSteps);
	assert(betas != NULL);
	for(int i = 0; i < timeSteps; i++)
	{
		printf("%f ", betas[i]);
	}
	free(betas);
}

void TestSigmoidBetas()
{
	int timeSteps = 4;
	float *betas = SigmoidBetaSchedule(timeSteps);
	assert(betas != NULL);
	for(int i = 0; i < timeSteps; i++)
	{
		printf("%f ", betas[i]);
	}
	free(betas);
}
int main()
{
	TestLinearSpace();
	TestLinearBetas();
	TestQuadraticBetas();
	TestCosineBetas();
	TestSigmoidBetas();
	return 0;
}