Skip to content

depth_fm.flow

flow

Flow-matching noise schedule (copied verbatim from upstream DepthFM).

cosine_alpha_bar 

cosine_alpha_bar(t: float) -> float

ᾱ_t under the cosine schedule.

Source code in src/depth_fm/flow/noise.py 
21
22
23
def cosine_alpha_bar(t: float) -> float:
    """ᾱ_t under the cosine schedule."""
    return 1 / (1 + np.exp(-cosine_log_snr(t)))

q_sample 

q_sample(x_start: Tensor, t: int, noise: Tensor = None, n_diffusion_timesteps: int = 1000) -> torch.Tensor

Diffuse x_start for t steps: sample from q(x_t | x_0).

Matches dfm.py exactly: alpha_bar_t = cosine_alpha_bar(t / n_diffusion_timesteps) return sqrt(alpha_bar_t)x_start + sqrt(1-alpha_bar_t)noise

Parameters:

Name Type Description Default
x_start Tensor

latent to noise, shape (B, C, H, W).

 required
t int

integer diffusion step (e.g. model.noising_step = 400).

 required
noise Tensor

optional pre-drawn noise; drawn fresh if None.

 None
n_diffusion_timesteps int

total diffusion steps (default 1000).

 1000

Returns:

Type Description
Tensor

Noised latent, same shape as x_start.
Source code in src/depth_fm/flow/noise.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
@torch.no_grad()
def q_sample(
        x_start: torch.Tensor,
        t: int,
        noise: torch.Tensor = None,
        n_diffusion_timesteps: int = 1000,
) -> torch.Tensor:
    """
    Diffuse x_start for t steps: sample from q(x_t | x_0).

    Matches dfm.py exactly:
        alpha_bar_t = cosine_alpha_bar(t / n_diffusion_timesteps)
        return sqrt(alpha_bar_t)*x_start + sqrt(1-alpha_bar_t)*noise

    Args:
        x_start: latent to noise, shape (B, C, H, W).
        t: integer diffusion step (e.g. model.noising_step = 400).
        noise: optional pre-drawn noise; drawn fresh if None.
        n_diffusion_timesteps: total diffusion steps (default 1000).

    Returns:
        Noised latent, same shape as x_start.
    """
    if noise is None:
        noise = torch.randn_like(x_start)
    alpha_bar_t = cosine_alpha_bar(t / n_diffusion_timesteps)
    alpha_bar_t = torch.tensor(alpha_bar_t, device=x_start.device, dtype=x_start.dtype)
    return torch.sqrt(alpha_bar_t) * x_start + torch.sqrt(1 - alpha_bar_t) * noise

per_sample_min_max_normalization 

per_sample_min_max_normalization(x: Tensor) -> torch.Tensor

Normalize each sample in a batch independently to [0, 1].

Matches dfm.py exactly.

Source code in src/depth_fm/flow/noise.py 
60
61
62
63
64
65
66
67
68
69
70
71
def per_sample_min_max_normalization(x: torch.Tensor) -> torch.Tensor:
    """
    Normalize each sample in a batch independently to [0, 1].

    Matches dfm.py exactly.
    """
    bs, *shape = x.shape
    x_ = einops.rearrange(x, "b ... -> b (...)")
    min_val = einops.reduce(x_, "b ... -> b", "min")[..., None]
    max_val = einops.reduce(x_, "b ... -> b", "max")[..., None]
    x_ = (x_ - min_val) / (max_val - min_val)
    return x_.reshape(bs, *shape)