luboslenco
863 lines
31 KiB
C
863 lines
31 KiB
C
/*
|
|
* Iris Sampling Implementation
|
|
*
|
|
* Rectified Flow sampling for image generation.
|
|
* Uses Euler method for ODE integration.
|
|
*/
|
|
|
|
#include "iris.h"
|
|
#include "iris_kernels.h"
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <sys/time.h>
|
|
#include <time.h>
|
|
|
|
#ifdef USE_METAL
|
|
#include "iris_metal.h"
|
|
#endif
|
|
|
|
/* Timing utilities for performance analysis - use wall-clock time */
|
|
static double get_time_ms(void) {
|
|
struct timeval tv;
|
|
gettimeofday(&tv, NULL);
|
|
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
|
}
|
|
|
|
/* Cumulative timing for denoising breakdown */
|
|
double iris_timing_transformer_total = 0.0;
|
|
double iris_timing_transformer_double = 0.0;
|
|
double iris_timing_transformer_single = 0.0;
|
|
double iris_timing_transformer_final = 0.0;
|
|
|
|
void iris_reset_timing(void) {
|
|
iris_timing_transformer_total = 0.0;
|
|
iris_timing_transformer_double = 0.0;
|
|
iris_timing_transformer_single = 0.0;
|
|
iris_timing_transformer_final = 0.0;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Timestep Schedules
|
|
* ======================================================================== */
|
|
|
|
/*
|
|
* Linear timestep schedule from 1.0 to 0.0
|
|
* Returns array of num_steps+1 values: [1.0, ..., 0.0]
|
|
*/
|
|
float *iris_schedule_linear(int num_steps) {
|
|
float *schedule = (float *)malloc((num_steps + 1) * sizeof(float));
|
|
for (int i = 0; i <= num_steps; i++) {
|
|
schedule[i] = 1.0f - (float)i / (float)num_steps;
|
|
}
|
|
return schedule;
|
|
}
|
|
|
|
/*
|
|
* Power schedule: denser steps at the start (high noise), sparser at the end.
|
|
* schedule[i] = 1 - (i/n)^alpha
|
|
* alpha=1.0 is linear, alpha=2.0 is quadratic, etc.
|
|
*/
|
|
float *iris_schedule_power(int num_steps, float alpha) {
|
|
float *schedule = (float *)malloc((num_steps + 1) * sizeof(float));
|
|
for (int i = 0; i <= num_steps; i++) {
|
|
float t = (float)i / (float)num_steps;
|
|
schedule[i] = 1.0f - powf(t, alpha);
|
|
}
|
|
return schedule;
|
|
}
|
|
|
|
/*
|
|
* Shifted sigmoid schedule (better for flow matching)
|
|
* shift controls where the inflection point is
|
|
*/
|
|
float *iris_schedule_sigmoid(int num_steps, float shift) {
|
|
float *schedule = (float *)malloc((num_steps + 1) * sizeof(float));
|
|
|
|
for (int i = 0; i <= num_steps; i++) {
|
|
float t = (float)i / (float)num_steps;
|
|
/* Shifted sigmoid: more steps at the end */
|
|
float x = (t - 0.5f) * 10.0f + shift;
|
|
schedule[i] = 1.0f - 1.0f / (1.0f + expf(-x));
|
|
}
|
|
|
|
/* Ensure endpoints */
|
|
schedule[0] = 1.0f;
|
|
schedule[num_steps] = 0.0f;
|
|
|
|
return schedule;
|
|
}
|
|
|
|
/*
|
|
* Resolution-dependent schedule (as used in FLUX.2)
|
|
* Higher resolutions use more steps at the start
|
|
*/
|
|
float *iris_schedule_resolution(int num_steps, int height, int width) {
|
|
float *schedule = (float *)malloc((num_steps + 1) * sizeof(float));
|
|
|
|
/* Compute shift based on resolution */
|
|
int pixels = height * width;
|
|
float shift = 0.0f;
|
|
if (pixels >= 1024 * 1024) {
|
|
shift = 1.0f; /* High res: more early steps */
|
|
}
|
|
else if (pixels >= 512 * 512) {
|
|
shift = 0.5f;
|
|
}
|
|
|
|
for (int i = 0; i <= num_steps; i++) {
|
|
float t = (float)i / (float)num_steps;
|
|
/* Apply shift */
|
|
t = powf(t, 1.0f + shift * 0.5f);
|
|
schedule[i] = 1.0f - t;
|
|
}
|
|
|
|
return schedule;
|
|
}
|
|
|
|
/*
|
|
* FLUX.2 official schedule with empirical mu calculation
|
|
* Matches Python's get_schedule() function from official flux2 code
|
|
*/
|
|
/* Compute the empirical shift parameter mu for the resolution-dependent
|
|
* noise schedule. The constants a1, b1, a2, b2 are fitted from the Flux
|
|
* training distribution and control how the SNR schedule adapts to different
|
|
* image resolutions. Higher resolution images need more denoising steps
|
|
* at high noise levels. Interpolates between two linear fits based on
|
|
* step count, with a cutoff at 4300 tokens. */
|
|
static float compute_empirical_mu(int image_seq_len, int num_steps) {
|
|
const float a1 = 8.73809524e-05f, b1 = 1.89833333f;
|
|
const float a2 = 0.00016927f, b2 = 0.45666666f;
|
|
|
|
if (image_seq_len > 4300) {
|
|
return a2 * image_seq_len + b2;
|
|
}
|
|
|
|
float m_200 = a2 * image_seq_len + b2;
|
|
float m_10 = a1 * image_seq_len + b1;
|
|
|
|
float a = (m_200 - m_10) / 190.0f;
|
|
float b = m_200 - 200.0f * a;
|
|
return a * num_steps + b;
|
|
}
|
|
|
|
/* Apply exponential SNR (Signal-to-Noise Ratio) shift to a timestep.
|
|
* Maps t in [0,1] through t/(t + (1-t)*exp(-mu)), shifting the schedule
|
|
* toward more time spent at higher noise levels. The boundary guards
|
|
* (t<=0, t>=1) prevent division by zero. */
|
|
static float generalized_time_snr_shift(float t, float mu, float sigma) {
|
|
/* t / (1 - t) with exp(mu) shift */
|
|
if (t <= 0.0f)
|
|
return 0.0f;
|
|
if (t >= 1.0f)
|
|
return 1.0f;
|
|
return expf(mu) / (expf(mu) + powf(1.0f / t - 1.0f, sigma));
|
|
}
|
|
|
|
float *iris_schedule_flux(int num_steps, int image_seq_len) {
|
|
float *schedule = (float *)malloc((num_steps + 1) * sizeof(float));
|
|
float mu = compute_empirical_mu(image_seq_len, num_steps);
|
|
|
|
for (int i = 0; i <= num_steps; i++) {
|
|
float t = 1.0f - (float)i / (float)num_steps; /* Linear from 1 to 0 */
|
|
schedule[i] = generalized_time_snr_shift(t, mu, 1.0f);
|
|
}
|
|
|
|
return schedule;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Euler Sampler for Rectified Flow
|
|
* ======================================================================== */
|
|
|
|
/*
|
|
* Single Euler step:
|
|
* z_next = z_t + (t_next - t_curr) * v(z_t, t_curr)
|
|
*
|
|
* Where v is the velocity predicted by the model.
|
|
* In rectified flow: v = (z_data - z_noise) at timestep t
|
|
*/
|
|
|
|
typedef struct iris_transformer_flux iris_transformer_flux_t;
|
|
typedef struct iris_vae iris_vae_t;
|
|
|
|
/* Free cached mmap weights after denoising */
|
|
extern void iris_transformer_free_mmap_cache_flux(iris_transformer_flux_t *tf);
|
|
|
|
/* Forward declarations */
|
|
extern float *iris_transformer_forward_flux(iris_transformer_flux_t *tf, const float *img_latent, int img_h, int img_w, const float *txt_emb, int txt_seq,
|
|
float timestep);
|
|
|
|
/* Forward declaration for in-context conditioning (img2img) */
|
|
extern float *iris_transformer_forward_refs_flux(iris_transformer_flux_t *tf, const float *img_latent, int img_h, int img_w, const float *ref_latent, int ref_h,
|
|
int ref_w, int t_offset, const float *txt_emb, int txt_seq, float timestep);
|
|
|
|
/* Forward declaration for multi-reference conditioning */
|
|
typedef struct {
|
|
const float *latent; /* Reference latent in NCHW format */
|
|
int h, w; /* Latent dimensions */
|
|
int t_offset; /* RoPE T coordinate (10, 20, 30, ...) */
|
|
} iris_ref_t;
|
|
|
|
extern float *iris_transformer_forward_multirefs_flux(iris_transformer_flux_t *tf, const float *img_latent, int img_h, int img_w, const iris_ref_t *refs,
|
|
int num_refs, const float *txt_emb, int txt_seq, float timestep);
|
|
|
|
/* VAE decode for step image callback */
|
|
extern iris_image *iris_vae_decode(iris_vae_t *vae, const float *latent, int batch, int latent_h, int latent_w);
|
|
extern void iris_image_free(iris_image *img);
|
|
|
|
/* Euler ODE sampler for Flux distilled models (txt2img). Iterates
|
|
* z_{n+1} = z_n + dt * v(z_n, t) where v is the velocity predicted by
|
|
* the transformer. The schedule provides sigma values from 1 to 0;
|
|
* the timestep passed to the transformer equals sigma, which it scales
|
|
* internally by 1000 for the sinusoidal embedding. */
|
|
float *iris_sample_euler_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const float *text_emb, int text_seq,
|
|
const float *schedule, int num_steps, void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder; /* Reserved for future use */
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
/* Working buffers */
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
float *v_cond = NULL;
|
|
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
/* Reset timing counters */
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr; /* Negative for denoising */
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
/* Notify step start */
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Predict velocity with conditioning */
|
|
v_cond = iris_transformer_forward_flux(tf, z_curr, h, w, text_emb, text_seq, t_curr);
|
|
|
|
/* Euler step: z_next = z_curr + dt * v */
|
|
iris_axpy(z_curr, dt, v_cond, latent_size);
|
|
|
|
free(v_cond);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback) {
|
|
progress_callback(step + 1, num_steps);
|
|
}
|
|
|
|
/* Step image callback - decode and display intermediate result */
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Print timing summary */
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown:\n");
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
if (iris_timing_transformer_double > 0 || iris_timing_transformer_single > 0) {
|
|
fprintf(stderr, " Transformer breakdown:\n");
|
|
fprintf(stderr, " Double blocks: %.1f ms (%.1f%%)\n", iris_timing_transformer_double,
|
|
100.0 * iris_timing_transformer_double / iris_timing_transformer_total);
|
|
fprintf(stderr, " Single blocks: %.1f ms (%.1f%%)\n", iris_timing_transformer_single,
|
|
100.0 * iris_timing_transformer_single / iris_timing_transformer_total);
|
|
fprintf(stderr, " Final layer: %.1f ms (%.1f%%)\n", iris_timing_transformer_final,
|
|
100.0 * iris_timing_transformer_final / iris_timing_transformer_total);
|
|
fprintf(stderr, " Total: %.1f ms\n", iris_timing_transformer_total);
|
|
/* Print fine-grained single block profile if available */
|
|
extern void iris_print_blas_profile(void);
|
|
iris_print_blas_profile();
|
|
}
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/* Euler sampler for single-reference img2img. The reference image is
|
|
* VAE-encoded and concatenated with the noised target as extra tokens.
|
|
* A RoPE T offset (default 10) distinguishes reference from target in
|
|
* the positional encoding. The transformer attends to both via joint
|
|
* attention, implementing in-context conditioning. */
|
|
float *iris_sample_euler_refs_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const float *ref_latent, int ref_h,
|
|
int ref_w, int t_offset, const float *text_emb, int text_seq, const float *schedule, int num_steps,
|
|
void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder; /* Reserved for future use */
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
/* Working buffer */
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
/* Reset timing counters */
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
/* Notify step start */
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Predict velocity with reference image conditioning */
|
|
float *v = iris_transformer_forward_refs_flux(tf, z_curr, h, w, ref_latent, ref_h, ref_w, t_offset, text_emb, text_seq, t_curr);
|
|
|
|
/* Euler step: z_next = z_curr + dt * v */
|
|
iris_axpy(z_curr, dt, v, latent_size);
|
|
|
|
free(v);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback) {
|
|
progress_callback(step + 1, num_steps);
|
|
}
|
|
|
|
/* Step image callback - decode and display intermediate result */
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Print timing summary */
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown (img2img with refs):\n");
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/*
|
|
* Sample using Euler method with multiple reference images.
|
|
* Each reference gets a different T offset in RoPE (10, 20, 30, ...).
|
|
*/
|
|
float *iris_sample_euler_multirefs_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const iris_ref_t *refs,
|
|
int num_refs, const float *text_emb, int text_seq, const float *schedule, int num_steps,
|
|
void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder;
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Predict velocity with multiple reference images */
|
|
float *v = iris_transformer_forward_multirefs_flux(tf, z_curr, h, w, refs, num_refs, text_emb, text_seq, t_curr);
|
|
|
|
/* Euler step */
|
|
iris_axpy(z_curr, dt, v, latent_size);
|
|
free(v);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback)
|
|
progress_callback(step + 1, num_steps);
|
|
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown (multi-ref, %d refs):\n", num_refs);
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* CFG (Classifier-Free Guidance) Samplers for Base Model
|
|
*
|
|
* These run the transformer twice per step: once with empty text (uncond)
|
|
* and once with the real prompt (cond), then combine:
|
|
* v = v_uncond + guidance_scale * (v_cond - v_uncond)
|
|
* ======================================================================== */
|
|
|
|
/* Euler sampler with Classifier-Free Guidance for Flux base models.
|
|
* Each step runs the transformer twice: once with empty prompt (unconditional)
|
|
* and once with real prompt (conditional). Combined as
|
|
* v = v_uncond + guidance_scale * (v_cond - v_uncond), which steers
|
|
* generation toward the prompt at the cost of 2x compute per step. */
|
|
float *iris_sample_euler_cfg_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const float *text_emb_cond,
|
|
int text_seq_cond, const float *text_emb_uncond, int text_seq_uncond, float guidance_scale, const float *schedule,
|
|
int num_steps, void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder;
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Unconditioned prediction */
|
|
float *v_uncond = iris_transformer_forward_flux(tf, z_curr, h, w, text_emb_uncond, text_seq_uncond, t_curr);
|
|
|
|
/* Conditioned prediction */
|
|
float *v_cond = iris_transformer_forward_flux(tf, z_curr, h, w, text_emb_cond, text_seq_cond, t_curr);
|
|
|
|
/* CFG combine: v = v_uncond + scale * (v_cond - v_uncond) */
|
|
for (int i = 0; i < latent_size; i++) {
|
|
float v = v_uncond[i] + guidance_scale * (v_cond[i] - v_uncond[i]);
|
|
z_curr[i] += dt * v;
|
|
}
|
|
|
|
free(v_uncond);
|
|
free(v_cond);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback)
|
|
progress_callback(step + 1, num_steps);
|
|
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown (CFG, guidance=%.1f):\n", guidance_scale);
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/*
|
|
* Euler sampler with CFG and single reference image (img2img).
|
|
*/
|
|
float *iris_sample_euler_cfg_refs_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const float *ref_latent,
|
|
int ref_h, int ref_w, int t_offset, const float *text_emb_cond, int text_seq_cond, const float *text_emb_uncond,
|
|
int text_seq_uncond, float guidance_scale, const float *schedule, int num_steps,
|
|
void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder;
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Unconditioned prediction (with ref) */
|
|
float *v_uncond = iris_transformer_forward_refs_flux(tf, z_curr, h, w, ref_latent, ref_h, ref_w, t_offset, text_emb_uncond, text_seq_uncond, t_curr);
|
|
|
|
/* Conditioned prediction (with ref) */
|
|
float *v_cond = iris_transformer_forward_refs_flux(tf, z_curr, h, w, ref_latent, ref_h, ref_w, t_offset, text_emb_cond, text_seq_cond, t_curr);
|
|
|
|
/* CFG combine */
|
|
for (int i = 0; i < latent_size; i++) {
|
|
float v = v_uncond[i] + guidance_scale * (v_cond[i] - v_uncond[i]);
|
|
z_curr[i] += dt * v;
|
|
}
|
|
|
|
free(v_uncond);
|
|
free(v_cond);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback)
|
|
progress_callback(step + 1, num_steps);
|
|
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown (CFG img2img, guidance=%.1f):\n", guidance_scale);
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/*
|
|
* Euler sampler with CFG and multiple reference images.
|
|
*/
|
|
float *iris_sample_euler_cfg_multirefs_flux(void *transformer, void *text_encoder, float *z, int batch, int channels, int h, int w, const iris_ref_t *refs,
|
|
int num_refs, const float *text_emb_cond, int text_seq_cond, const float *text_emb_uncond, int text_seq_uncond,
|
|
float guidance_scale, const float *schedule, int num_steps, void (*progress_callback)(int step, int total)) {
|
|
(void)text_encoder;
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
iris_reset_timing();
|
|
double total_denoising_start = get_time_ms();
|
|
double step_times[IRIS_MAX_STEPS];
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
double step_start = get_time_ms();
|
|
|
|
if (iris_step_callback)
|
|
iris_step_callback(step + 1, num_steps);
|
|
|
|
/* Unconditioned prediction (with refs) */
|
|
float *v_uncond = iris_transformer_forward_multirefs_flux(tf, z_curr, h, w, refs, num_refs, text_emb_uncond, text_seq_uncond, t_curr);
|
|
|
|
/* Conditioned prediction (with refs) */
|
|
float *v_cond = iris_transformer_forward_multirefs_flux(tf, z_curr, h, w, refs, num_refs, text_emb_cond, text_seq_cond, t_curr);
|
|
|
|
/* CFG combine */
|
|
for (int i = 0; i < latent_size; i++) {
|
|
float v = v_uncond[i] + guidance_scale * (v_cond[i] - v_uncond[i]);
|
|
z_curr[i] += dt * v;
|
|
}
|
|
|
|
free(v_uncond);
|
|
free(v_cond);
|
|
|
|
step_times[step] = get_time_ms() - step_start;
|
|
|
|
if (progress_callback)
|
|
progress_callback(step + 1, num_steps);
|
|
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (iris_verbose) {
|
|
double total_denoising = get_time_ms() - total_denoising_start;
|
|
fprintf(stderr, "\nDenoising timing breakdown (CFG multi-ref, %d refs, guidance=%.1f):\n", num_refs, guidance_scale);
|
|
for (int step = 0; step < num_steps; step++) {
|
|
fprintf(stderr, " Step %d: %.1f ms\n", step + 1, step_times[step]);
|
|
}
|
|
fprintf(stderr, " Total denoising: %.1f ms (%.2f s)\n", total_denoising, total_denoising / 1000.0);
|
|
}
|
|
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/*
|
|
* Sample using Euler method with stochastic noise injection.
|
|
* This can help with diversity and quality.
|
|
*/
|
|
float *iris_sample_euler_ancestral(void *transformer, float *z, int batch, int channels, int h, int w, const float *text_emb, int text_seq,
|
|
const float *schedule, int num_steps, float eta, void (*progress_callback)(int step, int total)) {
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
float *noise = (float *)malloc(latent_size * sizeof(float));
|
|
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
/* Predict velocity */
|
|
float *v = iris_transformer_forward_flux(tf, z_curr, h, w, text_emb, text_seq, t_curr);
|
|
|
|
/* Euler step */
|
|
iris_axpy(z_curr, dt, v, latent_size);
|
|
|
|
/* Add noise (ancestral sampling) */
|
|
if (eta > 0 && step < num_steps - 1) {
|
|
float sigma = eta * sqrtf(fabsf(dt));
|
|
iris_randn(noise, latent_size);
|
|
iris_axpy(z_curr, sigma, noise, latent_size);
|
|
}
|
|
|
|
free(v);
|
|
|
|
if (progress_callback) {
|
|
progress_callback(step + 1, num_steps);
|
|
}
|
|
|
|
/* Step image callback - decode and display intermediate result */
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
free(noise);
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Heun Sampler (2nd order)
|
|
* ======================================================================== */
|
|
|
|
/*
|
|
* Heun's method (improved Euler):
|
|
* 1. Predict: z_pred = z_t + dt * v(z_t, t)
|
|
* 2. Correct: z_next = z_t + dt/2 * (v(z_t, t) + v(z_pred, t+dt))
|
|
*/
|
|
float *iris_sample_heun(void *transformer, float *z, int batch, int channels, int h, int w, const float *text_emb, int text_seq, const float *schedule,
|
|
int num_steps, void (*progress_callback)(int step, int total)) {
|
|
iris_transformer_flux_t *tf = (iris_transformer_flux_t *)transformer;
|
|
int latent_size = batch * channels * h * w;
|
|
|
|
float *z_curr = (float *)malloc(latent_size * sizeof(float));
|
|
float *z_pred = (float *)malloc(latent_size * sizeof(float));
|
|
|
|
iris_copy(z_curr, z, latent_size);
|
|
|
|
for (int step = 0; step < num_steps; step++) {
|
|
float t_curr = schedule[step];
|
|
float t_next = schedule[step + 1];
|
|
float dt = t_next - t_curr;
|
|
|
|
/* First velocity estimate */
|
|
float *v1 = iris_transformer_forward_flux(tf, z_curr, h, w, text_emb, text_seq, t_curr);
|
|
|
|
/* Predict next state */
|
|
iris_copy(z_pred, z_curr, latent_size);
|
|
iris_axpy(z_pred, dt, v1, latent_size);
|
|
|
|
/* Second velocity estimate (only if not last step) */
|
|
if (step < num_steps - 1) {
|
|
float *v2 = iris_transformer_forward_flux(tf, z_pred, h, w, text_emb, text_seq, t_next);
|
|
|
|
/* Heun correction: z_next = z_curr + dt/2 * (v1 + v2) */
|
|
for (int i = 0; i < latent_size; i++) {
|
|
z_curr[i] += 0.5f * dt * (v1[i] + v2[i]);
|
|
}
|
|
|
|
free(v2);
|
|
}
|
|
else {
|
|
/* Last step: just use Euler */
|
|
iris_axpy(z_curr, dt, v1, latent_size);
|
|
}
|
|
|
|
free(v1);
|
|
|
|
if (progress_callback) {
|
|
progress_callback(step + 1, num_steps);
|
|
}
|
|
|
|
/* Step image callback - decode and display intermediate result */
|
|
if (iris_step_image_callback && iris_step_image_vae && step + 1 < num_steps) {
|
|
iris_image *img = iris_vae_decode((iris_vae_t *)iris_step_image_vae, z_curr, 1, h, w);
|
|
if (img) {
|
|
iris_step_image_callback(step + 1, num_steps, img);
|
|
iris_image_free(img);
|
|
}
|
|
}
|
|
}
|
|
|
|
free(z_pred);
|
|
iris_transformer_free_mmap_cache_flux(tf);
|
|
return z_curr;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Latent Noise Initialization
|
|
* ======================================================================== */
|
|
|
|
/* Generate initial noise for the denoising process. Uses a fixed 112x112
|
|
* noise patch tiled/cropped to the target size, ensuring seed-reproducible
|
|
* results regardless of output resolution. Without this, the same seed
|
|
* would produce different images at different sizes. For targets at or
|
|
* above 112x112 latent dims, noise is generated directly. */
|
|
#define NOISE_MAX_LATENT_DIM 112 /* 1792/16 = 112 */
|
|
|
|
float *iris_init_noise(int batch, int channels, int h, int w, int64_t seed) {
|
|
int target_size = batch * channels * h * w;
|
|
float *noise = (float *)malloc(target_size * sizeof(float));
|
|
|
|
if (seed >= 0) {
|
|
iris_rng_seed((uint64_t)seed);
|
|
}
|
|
|
|
/* If target is max size or larger, just generate directly */
|
|
if (h >= NOISE_MAX_LATENT_DIM && w >= NOISE_MAX_LATENT_DIM) {
|
|
iris_randn(noise, target_size);
|
|
return noise;
|
|
}
|
|
|
|
/* Generate noise at max latent size */
|
|
int max_h = NOISE_MAX_LATENT_DIM;
|
|
int max_w = NOISE_MAX_LATENT_DIM;
|
|
int max_size = batch * channels * max_h * max_w;
|
|
float *max_noise = (float *)malloc(max_size * sizeof(float));
|
|
iris_randn(max_noise, max_size);
|
|
|
|
/* Subsample to target size using nearest-neighbor */
|
|
for (int b = 0; b < batch; b++) {
|
|
for (int c = 0; c < channels; c++) {
|
|
for (int ty = 0; ty < h; ty++) {
|
|
for (int tx = 0; tx < w; tx++) {
|
|
/* Map target position to source position */
|
|
int sy = ty * max_h / h;
|
|
int sx = tx * max_w / w;
|
|
|
|
int src_idx = ((b * channels + c) * max_h + sy) * max_w + sx;
|
|
int dst_idx = ((b * channels + c) * h + ty) * w + tx;
|
|
noise[dst_idx] = max_noise[src_idx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
free(max_noise);
|
|
return noise;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Full Generation Pipeline
|
|
* ======================================================================== */
|
|
|
|
/* Complete text-to-image pipeline in latent space. Orchestrates:
|
|
* text encoding -> noise initialization -> schedule computation -> Euler
|
|
* sampling -> returns denoised latent (caller handles VAE decode).
|
|
* Routes to the appropriate sampler variant based on model type. */
|
|
typedef struct iris_ctx iris_ctx;
|
|
|
|
/* Forward declaration */
|
|
extern iris_ctx *iris_get_ctx(void);
|
|
|
|
float *iris_generate_latent(void *ctx_ptr, const float *text_emb, int text_seq, int height, int width, int num_steps, int64_t seed,
|
|
void (*progress_callback)(int step, int total)) {
|
|
/* Compute latent dimensions */
|
|
int latent_h = height / 16;
|
|
int latent_w = width / 16;
|
|
int channels = IRIS_LATENT_CHANNELS;
|
|
|
|
/* Initialize noise */
|
|
float *z = iris_init_noise(1, channels, latent_h, latent_w, seed);
|
|
|
|
/* Get schedule (4 steps for klein distilled) */
|
|
float *schedule = iris_schedule_linear(num_steps);
|
|
|
|
/* Sample (FLUX.2-klein is guidance-distilled, no CFG needed) */
|
|
float *latent = iris_sample_euler_flux(ctx_ptr, NULL, z, 1, channels, latent_h, latent_w, text_emb, text_seq, schedule, num_steps, progress_callback);
|
|
|
|
free(z);
|
|
free(schedule);
|
|
|
|
return latent;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Legacy Progress Callback (for backwards compatibility)
|
|
* ======================================================================== */
|
|
|
|
/* Legacy callback for step-level progress (called from sampling loop) */
|
|
void (*iris_progress_callback)(int, int) = NULL;
|
|
|
|
void iris_set_progress_callback(void (*callback)(int, int)) {
|
|
iris_progress_callback = callback;
|
|
}
|