class

DIAMONDModel

extendsPretrainedModel

DIAMONDModel(config: DIAMONDConfig)

source edit

The three networks that make up a DIAMOND agent.

Parameters

configDIAMONDConfig

Frozen configuration.

Attributes

denoiser_Denoiser

The U-Net whose preconditioned wrapper predicts the next frame.

reward_end_RewardEndModel

Reward class and termination, over a sequence.

actor_critic_ActorCritic

The policy and its value baseline.

Notes

Reference: Alonso, Eloi, et al., "Diffusion for World Modeling: Visual Details Matter in Atari", NeurIPS, 2024 (arXiv:2405.12399).

The EDM preconditioners live on this class rather than inside the U-Net, because they are a property of the diffusion, not of the network: the same U-Net under DDPM's parameterisation is the comparison the paper's Section 5.1 runs, and it drifts.

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config).eval()
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> with lucid.no_grad():
...     nxt = model.imagine_frame(frames, actions)
>>> nxt.shape
(2, 3, 16, 16)

Used by 2

Constructors

dunder

init

→None

__init__(config: DIAMONDConfig)

source edit

Initialise the agent. See the class docstring for parameters.

Instance methods

denoise

→Tensor

denoise(noised: Tensor, sigma: Tensor, frames: Tensor, actions: Tensor, cond_sigma: Tensor | None = None, quantize: bool = True)

source edit

Apply $D_\theta$ — the preconditioned denoiser.

Parameters

noisedTensor

(B, C, H, W), the next frame with noise added.

sigmaTensor

(B,) noise levels.

framesTensor

(B, L, C, H, W) clean history.

actionsTensor

(B, L) indices, or (B, L, num_actions) multi-hot.

cond_sigmaTensor or None= None

(B,) level to noise the history at. Only meaningful when the configuration asks for it; None leaves the history clean, which is what Atari does.

quantizebool= True, keyword-only

Put the estimate back on the 8-bit grid, as the released sampler does. Training passes False: the quantiser has no gradient, and the loss belongs on the estimate anyway.

Returns

Tensor

(B, C, H, W) estimate of the clean next frame.

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config).eval()
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> noised, sigma = lucid.randn((2, 3, 16, 16)), lucid.tensor([0.5, 2.0])
>>> clean = model.denoise(noised, sigma, frames, actions)
>>> clean.shape
(2, 3, 16, 16)
By default the estimate is put back on the 8-bit grid of [-1, 1],
which has no gradient; training asks for the raw estimate instead:
>>> bool(clean.min() >= -1.0), bool(clean.max() <= 1.0), clean.requires_grad
(True, True, False)
>>> model.denoise(noised, sigma, frames, actions, quantize=False).requires_grad
True

forward

→DIAMONDOutput

forward(frames: Tensor, actions: Tensor, next_frame: Tensor, sigma: Tensor | None = None)

source edit

Train the denoiser on one transition.

Parameters

framesTensor

(B, L, C, H, W) clean history.

actionsTensor

(B, L) indices, or (B, L, num_actions) multi-hot.

next_frameTensor

(B, C, H, W) the frame to predict.

sigma(Tensor or None, optional, keyword - only)= None

(B,) noise levels. Drawn from the paper's log-normal when absent; supplied by tests that need a fixed level.

Returns

DIAMONDOutput

Loss, denoised frame, and the noise levels used.

Notes

Reference: Alonso et al., arXiv:2405.12399, Algorithm 1 — sample $\log \sigma \sim \mathcal{N}(P_{\text{mean}}, P_{\text{std}}^2)$ , noise the target, and take the squared error in pixel space. The EDM weighting is already inside $D_\theta$ , so no extra loss weight appears here.

imagine_frame

→Tensor

imagine_frame(frames: Tensor, actions: Tensor, steps: int | None = None, noise: Tensor | None = None, cond_sigma: float | None = None)

source edit

Sample the next frame with Euler's method.

Parameters

framesTensor

(B, L, C, H, W) clean history.

actionsTensor

(B, L) indices, or (B, L, num_actions) multi-hot.

stepsint or None= None

Denoising steps; defaults to the configured 3.

cond_sigma(float or None, optional, keyword - only)= None

Noise the history at this level, and tell the network so. The released CS:GO sampler uses 0.005; None leaves the history clean, which is what Atari does.

noiseTensor or None= None

Starting sample, drawn when absent.

Returns

Tensor

(B, C, H, W).

Notes

Each step moves along $\mathrm{d}x/\mathrm{d}\sigma = (x - D_\theta(x, \sigma)) / \sigma$ , which is the probability-flow ODE written in EDM's variables.

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config).eval()
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> noise = lucid.randn((2, 3, 16, 16))
>>> with lucid.no_grad():
...     first = model.imagine_frame(frames, actions, noise=noise)
...     again = model.imagine_frame(frames, actions, noise=noise)
...     single = model.imagine_frame(frames, actions, noise=noise, steps=1)
>>> first.shape
(2, 3, 16, 16)
The sampler solves an ODE, so the starting noise fixes the frame and
a rollout can be reproduced; the number of steps changes it:
>>> bool((first == again).all()), bool((first == single).all())
(True, False)

preconditioners

→(Tensor, Tensor, Tensor, Tensor)

preconditioners(sigma: Tensor)

source edit

EDM's four scalings at a noise level.

Parameters

sigmaTensor

(B,) noise levels.

Returns

(Tensor, Tensor, Tensor, Tensor)

c_in, c_out, c_skip — each (B, 1, 1, 1) so they broadcast over a frame — and c_noise as (B,).

Notes

Reference: Alonso et al., arXiv:2405.12399, Appendix C, equations 9-12, with $\sigma_{\text{data}} = 0.5$ .

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config)
>>> c_in, c_out, c_skip, c_noise = model.preconditioners(
...     lucid.tensor([0.002, 0.5, 5.0]))
>>> c_in.shape, c_noise.shape
((3, 1, 1, 1), (3,))
The skip hands over to the network as the noise grows, until the
network is asked for the clean frame rather than for the noise:
>>> [round(v, 3) for v in c_skip.reshape(-1).tolist()]
[0.735, 0.424, 0.01]
The low end stops short of 1 because the offset noise is folded
into every level first, so even a clean level is scaled as if it
carried sigma_offset_noise:
>>> data, offset = config.sigma_data, config.sigma_offset_noise
>>> round(data**2 / (data**2 + offset**2), 3)
0.735

sigma_schedule

→Tensor

sigma_schedule(steps: int, device: str)

source edit

The noise levels an Euler sampler walks down.

Parameters

stepsint

Number of denoising steps.

devicestr

Where to build the schedule.

Returns

Tensor

(steps + 1,) descending to exactly zero.

Notes

Karras et al.'s $\rho$ -schedule, which Algorithm 1 refers to as "the default identity schedule from EDM". ⚠️ DIAMOND's tables give the sampler and the step count but not $\sigma_{\min}$ , $\sigma_{\max}$ or $\rho$ , so those come from EDM itself and are the one place here where a number is inherited rather than cited.

Examples

>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config)
>>> schedule = model.sigma_schedule(3, "cpu")
>>> schedule.shape
(4,)
EDM's sigma_max down to its sigma_min, then an exact zero, so
the last Euler step lands on the denoiser's own estimate:
>>> [round(v, 4) for v in schedule.tolist()]
[5.0, 0.2831, 0.002, 0.0]

step_actor_critic

→(Tensor, Tensor, tuple of Tensor)

step_actor_critic(frame: Tensor, state: tuple[Tensor, Tensor] | None = None)

source edit

One step of the actor-critic, typed.

Parameters

frameTensor

(B, C, H, W).

statetuple of Tensor or None= None

Carried LSTM state.

Returns

(Tensor, Tensor, tuple of Tensor)

Action logits, state value, and the new state.

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config).eval()
>>> frame = lucid.randn((2, 3, 16, 16))
>>> with lucid.no_grad():
...     logits, value, state = model.step_actor_critic(frame)
...     carried, _, _ = model.step_actor_critic(frame, state)
>>> logits.shape, value.shape, state[0].shape
((2, 4), (2,), (2, 16))
The policy has a memory: the same frame read with the carried state
gives different logits than it did from a blank one.
>>> bool((carried == logits).all())
False

step_reward_end

→(Tensor, Tensor, tuple of Tensor)

step_reward_end(frame: Tensor, next_frame: Tensor, action: Tensor, state: tuple[Tensor, Tensor] | None = None)

source edit

One step of the reward/termination model, typed.

Parameters

frameTensor

(B, C, H, W) each — the transition. Both are needed because a reward is a property of the transition, and the released encoder reads them stacked.

next_frameTensor

(B, C, H, W) each — the transition. Both are needed because a reward is a property of the transition, and the released encoder reads them stacked.

actionTensor

(B,) action indices taken from frame.

statetuple of Tensor or None= None

Carried LSTM state.

Returns

(Tensor, Tensor, tuple of Tensor)

Reward class logits (B, 3), termination class logits (B, 2), and the new state.

Examples

>>> import lucid
>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel
>>> config = DIAMONDConfig(
...     sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
...     reward_channels=(8, 8), reward_layers=(1, 1),
...     actor_channels=(8, 8), actor_layers=(1, 1),
...     cond_dim=16, reward_cond_dim=8,
...     reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4)
>>> model = DIAMONDModel(config).eval()
>>> frame, after = lucid.randn((2, 3, 16, 16)), lucid.randn((2, 3, 16, 16))
>>> action = lucid.tensor([0, 3], dtype=lucid.int64)
>>> with lucid.no_grad():
...     reward, end, state = model.step_reward_end(frame, after, action)
>>> reward.shape, end.shape, [s.shape for s in state]
((2, 3), (2, 2), [(2, 16), (2, 16)])
The three reward classes are the clipped signs -1, 0, +1, in
that order, so their probabilities give an expected reward in
[-1, 1]:
>>> probs = lucid.softmax(reward, dim=-1)
>>> expected = probs[:, 2] - probs[:, 0]
>>> bool((expected.abs() <= 1).all())
True

upsample

→Tensor

upsample(noised: Tensor, sigma: Tensor, low_res: Tensor, previous: Tensor, quantize: bool = True)

source edit

Apply the upsampler's $D_\theta$ .

Parameters

noisedTensor

(B, C, H*f, W*f) the full-resolution frame with noise.

sigmaTensor

(B,) noise levels.

low_resTensor

(B, C, H, W) the frame the world model generated.

previousTensor

(B, C, H*f, W*f) the previous full-resolution frame, so the sharpening is temporally consistent rather than independent per frame.

quantizebool= True, keyword-only

Put the estimate back on the 8-bit grid, as the released sampler does. Pass False to keep a gradient — the quantiser has none.

Returns

Tensor

(B, C, H*f, W*f).

Raises

ValueError

If this configuration has no upsampler.

Examples

>>> import lucid
>>> from lucid.models import diamond_csgo
>>> model = diamond_csgo(
...     unet_channels=(8, 8), unet_layers=(1, 1), cond_dim=16,
...     attn_depths=(0, 0), upsampler_channels=(4, 4),
...     upsampler_layers=(1, 1), upsampler_attn_depths=(0, 0),
...     num_actions=4).eval()
>>> low = lucid.randn((1, 3, 30, 56))
>>> noised = lucid.randn((1, 3, 150, 280))
>>> previous = lucid.randn((1, 3, 150, 280))
>>> with lucid.no_grad():
...     full = model.upsample(noised, lucid.tensor([1.0]), low, previous)
>>> full.shape
(1, 3, 150, 280)
One denoiser evaluation at one noise level — upsample_frame
walks it down a schedule — quantised back onto [-1, 1]:
>>> bool(full.abs().max() <= 1.0)
True

upsample_frame

→Tensor

upsample_frame(low_res: Tensor, previous: Tensor | None = None, steps: int | None = None, noise: Tensor | None = None)

source edit

Sample a full-resolution frame from a generated low-resolution one.

Parameters

low_resTensor

(B, C, H, W) — what imagine_frame produced.

previousTensor or None= None

(B, C, H*f, W*f) previous full-resolution frame. Absent on the first frame of a rollout, where the low-resolution one scaled up stands in for it.

stepsint or None= None

Denoising steps; defaults to the configured 3.

noiseTensor or None= None

Starting sample, drawn when absent.

Returns

Tensor

(B, C, H*f, W*f).

Raises

ValueError

If this configuration has no upsampler, or steps is not positive.

Notes

⚠️ The upsampler appears in the released CS:GO configuration and nowhere in the paper. It is why that experiment is affordable: the world model diffuses at 30x56 and this brings the frame to 150x280, rather than paying full-resolution diffusion for detail a cheaper network can add.

Examples

>>> import lucid
>>> from lucid.models import diamond_csgo
>>> model = diamond_csgo(
...     unet_channels=(8, 8), unet_layers=(1, 1), cond_dim=16,
...     attn_depths=(0, 0), upsampler_channels=(4, 4),
...     upsampler_layers=(1, 1), upsampler_attn_depths=(0, 0),
...     num_actions=4).eval()
>>> low = lucid.randn((1, 3, 30, 56))
>>> with lucid.no_grad():
...     full = model.upsample_frame(low, steps=1)
>>> full.shape
(1, 3, 150, 280)

>>> import lucid >>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config).eval() >>> frames = lucid.randn((2, 4, 3, 16, 16)) >>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64) >>> noised, sigma = lucid.randn((2, 3, 16, 16)), lucid.tensor([0.5, 2.0]) >>> clean = model.denoise(noised, sigma, frames, actions) >>> clean.shape (2, 3, 16, 16) By default the estimate is put back on the 8-bit grid of [-1, 1], which has no gradient; training asks for the raw estimate instead: >>> bool(clean.min() >= -1.0), bool(clean.max() <= 1.0), clean.requires_grad (True, True, False) >>> model.denoise(noised, sigma, frames, actions, quantize=False).requires_grad True

forward

→DIAMONDOutput

forward(frames: Tensor, actions: Tensor, next_frame: Tensor, sigma: Tensor | None = None)

source edit

Train the denoiser on one transition.

Parameters

framesTensor

(B, L, C, H, W) clean history.

actionsTensor

(B, L) indices, or (B, L, num_actions) multi-hot.

next_frameTensor

(B, C, H, W) the frame to predict.

sigma(Tensor or None, optional, keyword - only)= None

(B,) noise levels. Drawn from the paper's log-normal when absent; supplied by tests that need a fixed level.

Returns

DIAMONDOutput

Loss, denoised frame, and the noise levels used.

Notes

>>> import lucid >>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config).eval() >>> frames = lucid.randn((2, 4, 3, 16, 16)) >>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64) >>> noise = lucid.randn((2, 3, 16, 16)) >>> with lucid.no_grad(): ... first = model.imagine_frame(frames, actions, noise=noise) ... again = model.imagine_frame(frames, actions, noise=noise) ... single = model.imagine_frame(frames, actions, noise=noise, steps=1) >>> first.shape (2, 3, 16, 16) The sampler solves an ODE, so the starting noise fixes the frame and a rollout can be reproduced; the number of steps changes it: >>> bool((first == again).all()), bool((first == single).all()) (True, False)

>>> import lucid >>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config) >>> c_in, c_out, c_skip, c_noise = model.preconditioners( ... lucid.tensor([0.002, 0.5, 5.0])) >>> c_in.shape, c_noise.shape ((3, 1, 1, 1), (3,)) The skip hands over to the network as the noise grows, until the network is asked for the clean frame rather than for the noise: >>> [round(v, 3) for v in c_skip.reshape(-1).tolist()] [0.735, 0.424, 0.01] The low end stops short of 1 because the offset noise is folded into every level first, so even a clean level is scaled as if it carried sigma_offset_noise: >>> data, offset = config.sigma_data, config.sigma_offset_noise >>> round(data**2 / (data**2 + offset**2), 3) 0.735

>>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config) >>> schedule = model.sigma_schedule(3, "cpu") >>> schedule.shape (4,) EDM's sigma_max down to its sigma_min, then an exact zero, so the last Euler step lands on the denoiser's own estimate: >>> [round(v, 4) for v in schedule.tolist()] [5.0, 0.2831, 0.002, 0.0]

>>> import lucid >>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config).eval() >>> frame = lucid.randn((2, 3, 16, 16)) >>> with lucid.no_grad(): ... logits, value, state = model.step_actor_critic(frame) ... carried, _, _ = model.step_actor_critic(frame, state) >>> logits.shape, value.shape, state[0].shape ((2, 4), (2,), (2, 16)) The policy has a memory: the same frame read with the carried state gives different logits than it did from a blank one. >>> bool((carried == logits).all()) False

>>> import lucid >>> from lucid.models.generative.diamond import DIAMONDConfig, DIAMONDModel >>> config = DIAMONDConfig( ... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1), ... reward_channels=(8, 8), reward_layers=(1, 1), ... actor_channels=(8, 8), actor_layers=(1, 1), ... cond_dim=16, reward_cond_dim=8, ... reward_lstm_dim=16, actor_lstm_dim=16, num_actions=4) >>> model = DIAMONDModel(config).eval() >>> frame, after = lucid.randn((2, 3, 16, 16)), lucid.randn((2, 3, 16, 16)) >>> action = lucid.tensor([0, 3], dtype=lucid.int64) >>> with lucid.no_grad(): ... reward, end, state = model.step_reward_end(frame, after, action) >>> reward.shape, end.shape, [s.shape for s in state] ((2, 3), (2, 2), [(2, 16), (2, 16)]) The three reward classes are the clipped signs -1, 0, +1, in that order, so their probabilities give an expected reward in [-1, 1]: >>> probs = lucid.softmax(reward, dim=-1) >>> expected = probs[:, 2] - probs[:, 0] >>> bool((expected.abs() <= 1).all()) True

>>> import lucid >>> from lucid.models import diamond_csgo >>> model = diamond_csgo( ... unet_channels=(8, 8), unet_layers=(1, 1), cond_dim=16, ... attn_depths=(0, 0), upsampler_channels=(4, 4), ... upsampler_layers=(1, 1), upsampler_attn_depths=(0, 0), ... num_actions=4).eval() >>> low = lucid.randn((1, 3, 30, 56)) >>> noised = lucid.randn((1, 3, 150, 280)) >>> previous = lucid.randn((1, 3, 150, 280)) >>> with lucid.no_grad(): ... full = model.upsample(noised, lucid.tensor([1.0]), low, previous) >>> full.shape (1, 3, 150, 280) One denoiser evaluation at one noise level — upsample_frame walks it down a schedule — quantised back onto [-1, 1]: >>> bool(full.abs().max() <= 1.0) True

>>> import lucid >>> from lucid.models import diamond_csgo >>> model = diamond_csgo( ... unet_channels=(8, 8), unet_layers=(1, 1), cond_dim=16, ... attn_depths=(0, 0), upsampler_channels=(4, 4), ... upsampler_layers=(1, 1), upsampler_attn_depths=(0, 0), ... num_actions=4).eval() >>> low = lucid.randn((1, 3, 30, 56)) >>> with lucid.no_grad(): ... full = model.upsample_frame(low, steps=1) >>> full.shape (1, 3, 150, 280)