flux-latent-channel-animation.py

# This generates, along with the init latent, a second noise latent.
# Each of these has 16 channels, and the animation below has 16 phases.
# During phase i, the amount of noise in channel i decreases from max to
# 0, while the amount of noise in channel i + 1 increases from 0 to max.
# 
# For a render function and video export, see
# https://github.com/rolux/flux-random-walk

prompt = "a horse riding an astronaut"
width, height = 1024, 1024
seed = 42
frames = 15     # frames per channel
n = 16          # number of channels
max_noise = 1.0 # max amount of noise, decrease if needed
g = torch.Generator().manual_seed(seed)
latents = torch.randn((n, height//8, width//8), generator=g)
noise = torch.randn((n, height//8, width//8), generator=g) * max_noise

for c in range(n):
    for f in range(frames):
        i = c * frames + f
        filename = f"{prompt},{seed}/{i:08d}.png"
        t = f / frames
        dec = (1 - t) * noise[c,...] + t * latents[c,...]
        inc = (1 - t) * latents[(c+1)%n,...] + t * noise[(c+1)%n,...]
        dec, inc = dec.unsqueeze(0), inc.unsqueeze(0)
        if c < n - 1:
            concat = (latents[:c,...], dec, inc, latents[c+2:,...])
        else:
            concat = (inc, latents[1:-1,...], dec)
        c_latents = torch.concatenate(concat, axis=0)
        render(filename, prompt, width, height, c_latents)

[c,...] is just [c], obviously ;-)