import math
import torch
import numpy as np
from torch import nn
from typing import Callable, Iterable, Union, Optional
from einops import rearrange, repeat

from comfy import model_management
from .kl import (
	Encoder, Decoder, Upsample, Normalize,
	AttnBlock, ResnetBlock, #MemoryEfficientAttnBlock, 
	DiagonalGaussianDistribution, nonlinearity, make_attn
)

class AutoencoderKL(nn.Module):
	def __init__(self, config):
		super().__init__()
		self.embed_dim = config["embed_dim"]
		self.encoder = Encoder(**config)
		self.decoder = VideoDecoder(**config)
		assert config["double_z"]
		# these aren't used here for some reason
		# self.quant_conv = torch.nn.Conv2d(2*config["z_channels"], 2*self.embed_dim, 1)
		# self.post_quant_conv = torch.nn.Conv2d(self.embed_dim, config["z_channels"], 1)

	def encode(self, x):
		## batched
		# n_samples = x.shape[0]
		# n_rounds = math.ceil(x.shape[0] / n_samples)
		# all_out = []
		# for n in range(n_rounds):
			# h = self.encoder(
				# x[n * n_samples : (n + 1) * n_samples]
			# )
			# moments = h # self.quant_conv(h)
			# posterior = DiagonalGaussianDistribution(moments)
			# all_out.append(posterior.sample())
		# z = torch.cat(all_out, dim=0)
		# return z

		## default
		h = self.encoder(x)
		moments = h # self.quant_conv(h)
		posterior = DiagonalGaussianDistribution(moments)
		return posterior.sample()


	def decode(self, z):
		## batched - seems the same as default?
		# n_samples = z.shape[0]
		# n_rounds = math.ceil(z.shape[0] / n_samples)
		# all_out = []
		# for n in range(n_rounds):
			# dec = self.decoder(
				# z[n * n_samples : (n + 1) * n_samples],
				# timesteps=len(z[n * n_samples : (n + 1) * n_samples]),
			# )
			# all_out.append(dec)
		# out = torch.cat(all_out, dim=0)

		## default
		out = self.decoder(
			z, timesteps=len(z)
		)
		return out

	def forward(self, input, sample_posterior=True):
		posterior = self.encode(input)
		if sample_posterior:
			z = posterior.sample()
		else:
			z = posterior.mode()
		dec = self.decode(z)
		return dec, posterior

class VideoDecoder(nn.Module):
	available_time_modes = ["all", "conv-only", "attn-only"]
	def __init__(
		self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
		attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
		resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
		attn_type="vanilla",
		video_kernel_size: Union[int, list] = 3, alpha: float = 0.0, merge_strategy: str = "learned", time_mode: str = "conv-only",
		**ignorekwargs
	):
		super().__init__()
		if use_linear_attn: attn_type = "linear"
		self.ch = ch
		self.temb_ch = 0
		self.num_resolutions = len(ch_mult)
		self.num_res_blocks = num_res_blocks
		self.resolution = resolution
		self.in_channels = in_channels
		self.give_pre_end = give_pre_end
		self.tanh_out = tanh_out

		self.video_kernel_size = video_kernel_size
		self.alpha = alpha
		self.merge_strategy = merge_strategy
		self.time_mode = time_mode
		assert (
			self.time_mode in self.available_time_modes
		), f"time_mode parameter has to be in {self.available_time_modes}"

		# compute in_ch_mult, block_in and curr_res at lowest res
		in_ch_mult = (1,)+tuple(ch_mult)
		block_in = ch*ch_mult[self.num_resolutions-1]
		curr_res = resolution // 2**(self.num_resolutions-1)
		self.z_shape = (1,z_channels,curr_res,curr_res)
		print("Working with z of shape {} = {} dimensions.".format(
			self.z_shape, np.prod(self.z_shape)))

		# z to block_in
		self.conv_in = torch.nn.Conv2d(
			z_channels,
			block_in,
			kernel_size=3,
			stride=1,
			padding=1
		)

		# middle
		self.mid = nn.Module()
		self.mid.block_1 = VideoResBlock(
			in_channels=block_in,
			out_channels=block_in,
			temb_channels=self.temb_ch,
			dropout=dropout,
			video_kernel_size=self.video_kernel_size,
			alpha=self.alpha,
			merge_strategy=self.merge_strategy,
		)
		self.mid.attn_1 = make_attn(
			block_in,
			attn_type=attn_type,
		)
		self.mid.block_2 = VideoResBlock(
			in_channels=block_in,
			out_channels=block_in,
			temb_channels=self.temb_ch,
			dropout=dropout,
			video_kernel_size=self.video_kernel_size,
			alpha=self.alpha,
			merge_strategy=self.merge_strategy,
		)

		# upsampling
		self.up = nn.ModuleList()
		for i_level in reversed(range(self.num_resolutions)):
			block = nn.ModuleList()
			attn = nn.ModuleList()
			block_out = ch*ch_mult[i_level]
			for i_block in range(self.num_res_blocks+1):
				block.append(VideoResBlock(
					in_channels=block_in,
					out_channels=block_out,
					temb_channels=self.temb_ch,
					dropout=dropout,
					video_kernel_size=self.video_kernel_size,
					alpha=self.alpha,
					merge_strategy=self.merge_strategy,
				))
				block_in = block_out
				if curr_res in attn_resolutions:
					attn.append(make_time_attn(
						block_in,
						attn_type=attn_type,
						alpha=self.alpha,
						merge_strategy=self.merge_strategy,
					))
			up = nn.Module()
			up.block = block
			up.attn = attn
			if i_level != 0:
				up.upsample = Upsample(block_in, resamp_with_conv)
				curr_res = curr_res * 2
			self.up.insert(0, up) # prepend to get consistent order

		# end
		self.norm_out = Normalize(block_in)
		self.conv_out = AE3DConv(
			in_channels = block_in,
			out_channels = out_ch,
			video_kernel_size=self.video_kernel_size,
			kernel_size=3,
			stride=1,
			padding=1,
		)

	def get_last_layer(self, skip_time_mix=False, **kwargs):
		if self.time_mode == "attn-only":
			raise NotImplementedError("TODO")
		else:
			return (
				self.conv_out.time_mix_conv.weight
				if not skip_time_mix
				else self.conv_out.weight
			)

	def forward(self, z, **kwargs):
		#assert z.shape[1:] == self.z_shape[1:]
		self.last_z_shape = z.shape

		# timestep embedding
		temb = None

		# z to block_in
		h = self.conv_in(z)

		# middle
		h = self.mid.block_1(h, temb, **kwargs)
		h = self.mid.attn_1(h)
		h = self.mid.block_2(h, temb, **kwargs)

		# upsampling
		for i_level in reversed(range(self.num_resolutions)):
			for i_block in range(self.num_res_blocks+1):
				h = self.up[i_level].block[i_block](h, temb, **kwargs)
				if len(self.up[i_level].attn) > 0:
					h = self.up[i_level].attn[i_block](h, **kwargs)
			if i_level != 0:
				h = self.up[i_level].upsample(h)

		# end
		if self.give_pre_end:
			return h

		h = self.norm_out(h)
		h = nonlinearity(h)
		h = self.conv_out(h, **kwargs)
		if self.tanh_out:
			h = torch.tanh(h)
		return h

class CrossAttention(nn.Module):
	def __init__(
		self,
		query_dim,
		context_dim=None,
		heads=8,
		dim_head=64,
		dropout=0.0,
		backend=None,
	):
		super().__init__()
		inner_dim = dim_head * heads
		context_dim = context_dim or query_dim

		self.scale = dim_head**-0.5
		self.heads = heads

		self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
		self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
		self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

		self.to_out = nn.Sequential(
			nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
		)
		self.backend = backend

	def forward(
		self,
		x,
		context=None,
		mask=None,
		additional_tokens=None,
		n_times_crossframe_attn_in_self=0,
	):
		h = self.heads

		if additional_tokens is not None:
			# get the number of masked tokens at the beginning of the output sequence
			n_tokens_to_mask = additional_tokens.shape[1]
			# add additional token
			x = torch.cat([additional_tokens, x], dim=1)

		q = self.to_q(x)
		context = context or x
		k = self.to_k(context)
		v = self.to_v(context)

		if n_times_crossframe_attn_in_self:
			# reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
			assert x.shape[0] % n_times_crossframe_attn_in_self == 0
			n_cp = x.shape[0] // n_times_crossframe_attn_in_self
			k = repeat(
				k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
			)
			v = repeat(
				v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
			)

		q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))

		## old
		"""
		sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
		del q, k

		if exists(mask):
			mask = rearrange(mask, 'b ... -> b (...)')
			max_neg_value = -torch.finfo(sim.dtype).max
			mask = repeat(mask, 'b j -> (b h) () j', h=h)
			sim.masked_fill_(~mask, max_neg_value)

		# some quote about attention, that I just about had enough of
		sim = sim.softmax(dim=-1)

		out = einsum('b i j, b j d -> b i d', sim, v)
		"""
		## new
		with sdp_kernel(**BACKEND_MAP[self.backend]):
			# print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape)
			out = F.scaled_dot_product_attention(
				q, k, v, attn_mask=mask
			)  # scale is dim_head ** -0.5 per default

		del q, k, v
		out = rearrange(out, "b h n d -> b n (h d)", h=h)

		if additional_tokens is not None:
			# remove additional token
			out = out[:, n_tokens_to_mask:]
		return self.to_out(out)

class VideoBlock(AttnBlock):
	def __init__(
		self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"
	):
		super().__init__(in_channels)
		# no context, single headed, as in base class
		self.time_mix_block = VideoTransformerBlock(
			dim=in_channels,
			n_heads=1,
			d_head=in_channels,
			checkpoint=False,
			ff_in=True,
			attn_mode="softmax",
		)

		time_embed_dim = self.in_channels * 4
		self.video_time_embed = torch.nn.Sequential(
			torch.nn.Linear(self.in_channels, time_embed_dim),
			torch.nn.SiLU(),
			torch.nn.Linear(time_embed_dim, self.in_channels),
		)

		self.merge_strategy = merge_strategy
		if self.merge_strategy == "fixed":
			self.register_buffer("mix_factor", torch.Tensor([alpha]))
		elif self.merge_strategy == "learned":
			self.register_parameter(
				"mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))
			)
		else:
			raise ValueError(f"unknown merge strategy {self.merge_strategy}")

	def forward(self, x, timesteps, skip_video=False):
		if skip_video:
			return super().forward(x)

		x_in = x
		x = self.attention(x)
		h, w = x.shape[2:]
		x = rearrange(x, "b c h w -> b (h w) c")

		x_mix = x
		num_frames = torch.arange(timesteps, device=x.device)
		num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)
		num_frames = rearrange(num_frames, "b t -> (b t)")
		t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)
		emb = self.video_time_embed(t_emb)  # b, n_channels
		emb = emb[:, None, :]
		x_mix = x_mix + emb

		alpha = self.get_alpha()
		x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
		x = alpha * x + (1.0 - alpha) * x_mix  # alpha merge

		x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
		x = self.proj_out(x)

		return x_in + x

	def attention(self, h_: torch.Tensor) -> torch.Tensor:
		h_ = self.norm(h_)
		q = self.q(h_)
		k = self.k(h_)
		v = self.v(h_)

		b, c, h, w = q.shape
		q, k, v = map(
			lambda x: rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v)
		)
		h_ = torch.nn.functional.scaled_dot_product_attention(
			q, k, v
		)  # scale is dim ** -0.5 per default
		# compute attention

		return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)

	def forward(self, x, **kwargs):
		h_ = x
		h_ = self.attention(h_)
		h_ = self.proj_out(h_)
		return x + h_

	def get_alpha(
		self,
	):
		if self.merge_strategy == "fixed":
			return self.mix_factor
		elif self.merge_strategy == "learned":
			return torch.sigmoid(self.mix_factor)
		else:
			raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")

class VideoTransformerBlock(nn.Module):
	ATTENTION_MODES = {
		"softmax": CrossAttention,
		# "softmax-xformers": MemoryEfficientCrossAttention,
	}

	def __init__(
		self,
		dim,
		n_heads,
		d_head,
		dropout=0.0,
		context_dim=None,
		gated_ff=True,
		checkpoint=True,
		timesteps=None,
		ff_in=False,
		inner_dim=None,
		attn_mode="softmax",
		disable_self_attn=False,
		disable_temporal_crossattention=False,
		switch_temporal_ca_to_sa=False,
	):
		super().__init__()

		attn_cls = self.ATTENTION_MODES.get(attn_mode, "softmax")

		self.ff_in = ff_in or inner_dim is not None
		if inner_dim is None:
			inner_dim = dim

		assert int(n_heads * d_head) == inner_dim

		self.is_res = inner_dim == dim

		if self.ff_in:
			self.norm_in = nn.LayerNorm(dim)
			self.ff_in = FeedForward(
				dim, dim_out=inner_dim, dropout=dropout, glu=gated_ff
			)

		self.timesteps = timesteps
		self.disable_self_attn = disable_self_attn
		if self.disable_self_attn:
			self.attn1 = attn_cls(
				query_dim=inner_dim,
				heads=n_heads,
				dim_head=d_head,
				context_dim=context_dim,
				dropout=dropout,
			)  # is a cross-attention
		else:
			self.attn1 = attn_cls(
				query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout
			)  # is a self-attention

		self.ff = FeedForward(inner_dim, dim_out=dim, dropout=dropout, glu=gated_ff)

		if disable_temporal_crossattention:
			if switch_temporal_ca_to_sa:
				raise ValueError
			else:
				self.attn2 = None
		else:
			self.norm2 = nn.LayerNorm(inner_dim)
			if switch_temporal_ca_to_sa:
				self.attn2 = attn_cls(
					query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout
				)  # is a self-attention
			else:
				self.attn2 = attn_cls(
					query_dim=inner_dim,
					context_dim=context_dim,
					heads=n_heads,
					dim_head=d_head,
					dropout=dropout,
				)  # is self-attn if context is none

		self.norm1 = nn.LayerNorm(inner_dim)
		self.norm3 = nn.LayerNorm(inner_dim)
		self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa

		self.checkpoint = checkpoint
		if self.checkpoint:
			print(f"{self.__class__.__name__} is using checkpointing")

	def forward(
		self, x: torch.Tensor, context: torch.Tensor = None, timesteps: int = None
	) -> torch.Tensor:
		if self.checkpoint:
			return checkpoint(self._forward, x, context, timesteps)
		else:
			return self._forward(x, context, timesteps=timesteps)

	def _forward(self, x, context=None, timesteps=None):
		assert self.timesteps or timesteps
		assert not (self.timesteps and timesteps) or self.timesteps == timesteps
		timesteps = self.timesteps or timesteps
		B, S, C = x.shape
		x = rearrange(x, "(b t) s c -> (b s) t c", t=timesteps)

		if self.ff_in:
			x_skip = x
			x = self.ff_in(self.norm_in(x))
			if self.is_res:
				x += x_skip

		if self.disable_self_attn:
			x = self.attn1(self.norm1(x), context=context) + x
		else:
			x = self.attn1(self.norm1(x)) + x

		if self.attn2 is not None:
			if self.switch_temporal_ca_to_sa:
				x = self.attn2(self.norm2(x)) + x
			else:
				x = self.attn2(self.norm2(x), context=context) + x
		x_skip = x
		x = self.ff(self.norm3(x))
		if self.is_res:
			x += x_skip

		x = rearrange(
			x, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps
		)
		return x

	def get_last_layer(self):
		return self.ff.net[-1].weight

class ResBlock(nn.Module):
	"""
	A residual block that can optionally change the number of channels.
	:param channels: the number of input channels.
	:param emb_channels: the number of timestep embedding channels.
	:param dropout: the rate of dropout.
	:param out_channels: if specified, the number of out channels.
	:param use_conv: if True and out_channels is specified, use a spatial
		convolution instead of a smaller 1x1 convolution to change the
		channels in the skip connection.
	:param dims: determines if the signal is 1D, 2D, or 3D.
	:param use_checkpoint: if True, use gradient checkpointing on this module.
	:param up: if True, use this block for upsampling.
	:param down: if True, use this block for downsampling.
	"""

	def __init__(
		self,
		channels: int,
		emb_channels: int,
		dropout: float,
		out_channels: Optional[int] = None,
		use_conv: bool = False,
		use_scale_shift_norm: bool = False,
		dims: int = 2,
		use_checkpoint: bool = False,
		up: bool = False,
		down: bool = False,
		kernel_size: int = 3,
		exchange_temb_dims: bool = False,
		skip_t_emb: bool = False,
	):
		super().__init__()
		self.channels = channels
		self.emb_channels = emb_channels
		self.dropout = dropout
		self.out_channels = out_channels or channels
		self.use_conv = use_conv
		self.use_checkpoint = use_checkpoint
		self.use_scale_shift_norm = use_scale_shift_norm
		self.exchange_temb_dims = exchange_temb_dims

		if isinstance(kernel_size, Iterable):
			padding = [k // 2 for k in kernel_size]
		else:
			padding = kernel_size // 2

		self.in_layers = nn.Sequential(
			normalization(channels),
			nn.SiLU(),
			conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding),
		)

		self.updown = up or down

		if up:
			self.h_upd = Upsample(channels, False, dims)
			self.x_upd = Upsample(channels, False, dims)
		elif down:
			self.h_upd = Downsample(channels, False, dims)
			self.x_upd = Downsample(channels, False, dims)
		else:
			self.h_upd = self.x_upd = nn.Identity()

		self.skip_t_emb = skip_t_emb
		self.emb_out_channels = (
			2 * self.out_channels if use_scale_shift_norm else self.out_channels
		)
		if self.skip_t_emb:
			print(f"Skipping timestep embedding in {self.__class__.__name__}")
			assert not self.use_scale_shift_norm
			self.emb_layers = None
			self.exchange_temb_dims = False
		else:
			self.emb_layers = nn.Sequential(
				nn.SiLU(),
				linear(
					emb_channels,
					self.emb_out_channels,
				),
			)

		self.out_layers = nn.Sequential(
			normalization(self.out_channels),
			nn.SiLU(),
			nn.Dropout(p=dropout),
			zero_module(
				conv_nd(
					dims,
					self.out_channels,
					self.out_channels,
					kernel_size,
					padding=padding,
				)
			),
		)

		if self.out_channels == channels:
			self.skip_connection = nn.Identity()
		elif use_conv:
			self.skip_connection = conv_nd(
				dims, channels, self.out_channels, kernel_size, padding=padding
			)
		else:
			self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)

	def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
		"""
		Apply the block to a Tensor, conditioned on a timestep embedding.
		:param x: an [N x C x ...] Tensor of features.
		:param emb: an [N x emb_channels] Tensor of timestep embeddings.
		:return: an [N x C x ...] Tensor of outputs.
		"""
		if self.use_checkpoint:
			return checkpoint(self._forward, x, emb)
		else:
			return self._forward(x, emb)

	def _forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
		if self.updown:
			in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
			h = in_rest(x)
			h = self.h_upd(h)
			x = self.x_upd(x)
			h = in_conv(h)
		else:
			h = self.in_layers(x)

		if self.skip_t_emb:
			emb_out = torch.zeros_like(h)
		else:
			emb_out = self.emb_layers(emb).type(h.dtype)
		while len(emb_out.shape) < len(h.shape):
			emb_out = emb_out[..., None]
		if self.use_scale_shift_norm:
			out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
			scale, shift = torch.chunk(emb_out, 2, dim=1)
			h = out_norm(h) * (1 + scale) + shift
			h = out_rest(h)
		else:
			if self.exchange_temb_dims:
				emb_out = rearrange(emb_out, "b t c ... -> b c t ...")
			h = h + emb_out
			h = self.out_layers(h)
		return self.skip_connection(x) + h

class VideoResBlock(ResnetBlock):
	def __init__(
		self,
		out_channels,
		*args,
		dropout=0.0,
		video_kernel_size=3,
		alpha=0.0,
		merge_strategy="learned",
		**kwargs,
	):
		super().__init__(out_channels=out_channels, dropout=dropout, *args, **kwargs)
		if video_kernel_size is None:
			video_kernel_size = [3, 1, 1]
		self.time_stack = ResBlock(
			channels=out_channels,
			emb_channels=0,
			dropout=dropout,
			dims=3,
			use_scale_shift_norm=False,
			use_conv=False,
			up=False,
			down=False,
			kernel_size=video_kernel_size,
			use_checkpoint=False,
			skip_t_emb=True,
		)

		self.merge_strategy = merge_strategy
		if self.merge_strategy == "fixed":
			self.register_buffer("mix_factor", torch.Tensor([alpha]))
		elif self.merge_strategy == "learned":
			self.register_parameter(
				"mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))
			)
		else:
			raise ValueError(f"unknown merge strategy {self.merge_strategy}")

	def get_alpha(self, bs):
		if self.merge_strategy == "fixed":
			return self.mix_factor
		elif self.merge_strategy == "learned":
			return torch.sigmoid(self.mix_factor)
		else:
			raise NotImplementedError()

	def forward(self, x, temb, skip_video=False, timesteps=None):
		if timesteps is None:
			timesteps = self.timesteps

		b, c, h, w = x.shape

		x = super().forward(x, temb)

		if not skip_video:
			x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)

			x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)

			x = self.time_stack(x, temb)

			alpha = self.get_alpha(bs=b // timesteps)
			x = alpha * x + (1.0 - alpha) * x_mix

			x = rearrange(x, "b c t h w -> (b t) c h w")
		return x

class AE3DConv(torch.nn.Conv2d):
	def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs):
		super().__init__(in_channels, out_channels, *args, **kwargs)
		if isinstance(video_kernel_size, Iterable):
			padding = [int(k // 2) for k in video_kernel_size]
		else:
			padding = int(video_kernel_size // 2)

		self.time_mix_conv = torch.nn.Conv3d(
			in_channels=out_channels,
			out_channels=out_channels,
			kernel_size=video_kernel_size,
			padding=padding,
		)

	def forward(self, input, timesteps, skip_video=False):
		x = super().forward(input)
		if skip_video:
			return x
		x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
		x = self.time_mix_conv(x)
		return rearrange(x, "b c t h w -> (b t) c h w")

def make_time_attn(in_channels, attn_type="vanilla", attn_kwargs=None, alpha: float = 0, merge_strategy: str = "learned"):
	if attn_type == "vanilla":
		assert attn_kwargs is None
		return VideoBlock(
			in_channels,
			alpha=alpha,
			merge_strategy=merge_strategy,
		)
	# lazy to add the xformers code
	else:
		return NotImplementedError()

def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
	"""
	Create sinusoidal timestep embeddings.
	:param timesteps: a 1-D Tensor of N indices, one per batch element.
					  These may be fractional.
	:param dim: the dimension of the output.
	:param max_period: controls the minimum frequency of the embeddings.
	:return: an [N x dim] Tensor of positional embeddings.
	"""
	if not repeat_only:
		half = dim // 2
		freqs = torch.exp(
			-math.log(max_period)
			* torch.arange(start=0, end=half, dtype=torch.float32)
			/ half
		).to(device=timesteps.device)
		args = timesteps[:, None].float() * freqs[None]
		embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
		if dim % 2:
			embedding = torch.cat(
				[embedding, torch.zeros_like(embedding[:, :1])], dim=-1
			)
	else:
		embedding = repeat(timesteps, "b -> b d", d=dim)
	return embedding

def normalization(channels):
	"""
	Make a standard normalization layer.
	:param channels: number of input channels.
	:return: an nn.Module for normalization.
	"""
	return GroupNorm32(32, channels)

class SiLU(nn.Module):
	def forward(self, x):
		return x * torch.sigmoid(x)

class GroupNorm32(nn.GroupNorm):
	def forward(self, x):
		return super().forward(x.float()).type(x.dtype)

def conv_nd(dims, *args, **kwargs):
	"""
	Create a 1D, 2D, or 3D convolution module.
	"""
	if dims == 1:
		return nn.Conv1d(*args, **kwargs)
	elif dims == 2:
		return nn.Conv2d(*args, **kwargs)
	elif dims == 3:
		return nn.Conv3d(*args, **kwargs)
	raise ValueError(f"unsupported dimensions: {dims}")

def zero_module(module):
	"""
	Zero out the parameters of a module and return it.
	"""
	for p in module.parameters():
		p.detach().zero_()
	return module

# feedforward
class GEGLU(nn.Module):
	def __init__(self, dim_in, dim_out):
		super().__init__()
		self.proj = nn.Linear(dim_in, dim_out * 2)

	def forward(self, x):
		x, gate = self.proj(x).chunk(2, dim=-1)
		return x * nn.functional.gelu(gate)

class FeedForward(nn.Module):
	def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
		super().__init__()
		inner_dim = int(dim * mult)
		dim_out = dim_out or dim
		project_in = (
			nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
			if not glu
			else GEGLU(dim, inner_dim)
		)

		self.net = nn.Sequential(
			project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
		)

	def forward(self, x):
		return self.net(x)