File size: 1,116 Bytes
c5e4c7b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
"""
IRIS: Iterative Refinement Image Synthesizer
=============================================

A mobile-first image generation architecture that combines:
- PDE-SSM spatial mixing (O(N log N), native 2D, no scanning)
- Weight-shared iterative refinement (6 blocks × R iterations)
- Structured latent canvas (DC-AE with coarse-to-fine channels)
- Tiny PixelShuffle decoder (0.1M params)
- Multi-Query Attention + 2D RoPE
- Flow matching training with logit-normal timestep sampling

Usage:
    from iris.model import IRIS
    from iris.configs import get_model_config
    from iris.flow_matching import flow_matching_loss, euler_sample
    
    model = IRIS(**get_model_config("iris-small"))
    
    # Training
    losses = flow_matching_loss(model, latents, text_embeddings)
    
    # Sampling
    images = euler_sample(model, noise, text_embeddings, num_steps=20)
"""

from .model import IRIS
from .configs import get_model_config, CONFIGS
from .flow_matching import flow_matching_loss, euler_sample

__version__ = "0.1.0"
__all__ = ["IRIS", "get_model_config", "CONFIGS", "flow_matching_loss", "euler_sample"]