diffusion-1d/sampling.py at main · merdan-9/diffusion-1d · GitHub

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
"""Inference utilities for generating trajectories from the diffusion model."""

from typing import Optional

import torch

from diffusion import Diffusion1D
from guidance import SineWaveGuidance

from utils import set_seed


def generate_sequences(
    diffusion: Diffusion1D,
    num_samples: int,
    seq_length: int,
    device: torch.device,
    seed: Optional[int] = None,
    guidance: Optional[SineWaveGuidance] = None,
) -> torch.Tensor:
    """Sample new sine-like sequences using the trained diffusion model."""
    set_seed(seed or 42)

    # Create shape tuple
    shape = (num_samples, seq_length)

    # Generate samples
    samples = diffusion.sample(shape, guidance=guidance)

    return samples