Add files using upload-large-folder tool

31ab957 verified about 1 month ago

7.15 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# GLIDE: https://github.com/openai/glide-text2im
# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
# --------------------------------------------------------

import torch
import torch.nn as nn
import numpy as np
import math
from .modules import FinalLayer, DiTBlock
from x_transformers.x_transformers import RotaryEmbedding



#################################################################################
#               Embedding Layers for Timesteps and Class Labels                 #
#################################################################################



class SinusPositionEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x, scale=1000):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb

class TimestepEmbedder(nn.Module):
    def __init__(self, dim, freq_embed_dim=256):
        super().__init__()
        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))

    def forward(self, timestep):
        time_hidden = self.time_embed(timestep)
        time_hidden = time_hidden.to(timestep.dtype)
        time = self.time_mlp(time_hidden)  # b d
        return time


class CondEmbedder(nn.Module):
    """
    Embeds class labels into vector representations. Also handles llm hidden dropout for classifier-free guidance.
    """
    def __init__(self, input_feature_size, hidden_size, dropout_prob):
        super().__init__()
        self.dropout_prob = dropout_prob
        self.cond_embedder = nn.Linear(input_feature_size, hidden_size)

    def cond_drop(self, llm_cond):
        """
        Drops llm hidden to enable classifier-free guidance.
        """
        bsz = llm_cond.shape[0]
        drop_latent_mask = torch.rand(bsz) < self.dropout_prob
        drop_latent_mask = drop_latent_mask.unsqueeze(-1).unsqueeze(-1).to(llm_cond.dtype).to(llm_cond.device)
        fake_latent = torch.zeros(llm_cond.shape).to(llm_cond.device)
        llm_cond  = drop_latent_mask * fake_latent + (1 - drop_latent_mask) * llm_cond

        return llm_cond

    def forward(self, llm_cond, train):
        use_dropout = self.dropout_prob > 0
        if train and use_dropout:
            llm_cond = self.cond_drop(llm_cond)

        llm_cond = self.cond_embedder(llm_cond)


        return llm_cond



class DiT(nn.Module):
    """
    Diffusion model with a Transformer backbone.
    """
    def __init__(
        self,
        in_channels=4,
        hidden_size=1024,
        depth=28,
        num_heads=16,
        mlp_ratio=4.0,
        llm_cond_dim=896,
        cfg_dropout_prob=0.1,
        **kwargs,
    ):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = in_channels
        self.num_heads = num_heads


        self.t_embedder = TimestepEmbedder(hidden_size)
        self.x_embedder = nn.Linear(in_channels, hidden_size)
        self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size, cfg_dropout_prob)
        self.hidden_size = hidden_size

        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)


        self.blocks = nn.ModuleList([
            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)
        ])
        self.final_layer = FinalLayer(hidden_size, self.out_channels)
        self.initialize_weights()


    def initialize_weights(self):
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
        self.apply(_basic_init)


        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
        w_x = self.x_embedder.weight.data
        nn.init.xavier_uniform_(w_x.view([w_x.shape[0], -1]))
        nn.init.constant_(self.x_embedder.bias, 0)

        # Initialize label embedding table:
        w_c = self.c_embedder.cond_embedder.weight.data
        nn.init.xavier_uniform_(w_c.view([w_c.shape[0], -1]))
        nn.init.constant_(self.c_embedder.cond_embedder.bias, 0)

        # Initialize timestep embedding MLP:
        nn.init.normal_(self.t_embedder.time_mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder.time_mlp[2].weight, std=0.02)


        # Zero-out output layers:
        nn.init.constant_(self.final_layer.linear.weight, 0)
        nn.init.constant_(self.final_layer.linear.bias, 0)



    def forward(self, x, t, c, latent_history, mask=None):
        """
        Forward pass of DiT.
        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
        t: (N,) tensor of diffusion timesteps
        y: (N,) tensor of class labels
        """
        t = self.t_embedder(t).unsqueeze(1)          # (N, D)
        x_now = self.x_embedder(x)  #转换x_noise的维度
        x_history = self.x_embedder(latent_history) #转换target history的维度
        x = torch.cat([x_history, x_now], dim=1) # 拼接patch history和当前patch
        c = self.c_embedder(c, self.training)  # (N, 1, 896) -> (N, 1, D)，同时带有CFG替换操作
        y = t + c
        x = torch.cat([y, x], dim=1) # # (N, 1 + patch_size *2, D)
        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])

        if mask is not None:
            mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1])
            mask = torch.cat([mask_pad, mask], dim=-1)
        for block in self.blocks:
            x = block(x, mask, rope)                      # (N, T, D)
        x = self.final_layer(x)               # (N, T, patch_size ** 2 * out_channels)
        #x = self.unpatchify(x)                   # (N, out_channels, H, W)
        return x

    def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
        """
        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
        """
        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
        if not cfg_scale == 1:
            x = torch.cat([x, x], dim=0)
            latent_history = torch.cat([latent_history, latent_history], dim=0)
            fake_latent = torch.zeros(c.shape).to(c.device)
            c = torch.cat([c, fake_latent], dim=0)
        if t.ndim == 0:
            t = t.repeat(x.shape[0])
        model_out = self.forward(x, t, c, latent_history)
        return model_out[:, -patch_size:, :]

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.

	# This source code is licensed under the license found in the
	# LICENSE file in the root directory of this source tree.
	# --------------------------------------------------------
	# References:
	# GLIDE: https://github.com/openai/glide-text2im
	# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
	# --------------------------------------------------------

	import torch
	import torch.nn as nn
	import numpy as np
	import math
	from .modules import FinalLayer, DiTBlock
	from x_transformers.x_transformers import RotaryEmbedding



	#################################################################################
	# Embedding Layers for Timesteps and Class Labels #
	#################################################################################



	class SinusPositionEmbedding(nn.Module):
	def __init__(self, dim):
	super().__init__()
	self.dim = dim

	def forward(self, x, scale=1000):
	device = x.device
	half_dim = self.dim // 2
	emb = math.log(10000) / (half_dim - 1)
	emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
	emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
	emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
	return emb

	class TimestepEmbedder(nn.Module):
	def __init__(self, dim, freq_embed_dim=256):
	super().__init__()
	self.time_embed = SinusPositionEmbedding(freq_embed_dim)
	self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))

	def forward(self, timestep):
	time_hidden = self.time_embed(timestep)
	time_hidden = time_hidden.to(timestep.dtype)
	time = self.time_mlp(time_hidden) # b d
	return time


	class CondEmbedder(nn.Module):
	"""
	Embeds class labels into vector representations. Also handles llm hidden dropout for classifier-free guidance.
	"""
	def __init__(self, input_feature_size, hidden_size, dropout_prob):
	super().__init__()
	self.dropout_prob = dropout_prob
	self.cond_embedder = nn.Linear(input_feature_size, hidden_size)

	def cond_drop(self, llm_cond):
	"""
	Drops llm hidden to enable classifier-free guidance.
	"""
	bsz = llm_cond.shape[0]
	drop_latent_mask = torch.rand(bsz) < self.dropout_prob
	drop_latent_mask = drop_latent_mask.unsqueeze(-1).unsqueeze(-1).to(llm_cond.dtype).to(llm_cond.device)
	fake_latent = torch.zeros(llm_cond.shape).to(llm_cond.device)
	llm_cond = drop_latent_mask * fake_latent + (1 - drop_latent_mask) * llm_cond

	return llm_cond

	def forward(self, llm_cond, train):
	use_dropout = self.dropout_prob > 0
	if train and use_dropout:
	llm_cond = self.cond_drop(llm_cond)

	llm_cond = self.cond_embedder(llm_cond)


	return llm_cond



	class DiT(nn.Module):
	"""
	Diffusion model with a Transformer backbone.
	"""
	def __init__(
	self,
	in_channels=4,
	hidden_size=1024,
	depth=28,
	num_heads=16,
	mlp_ratio=4.0,
	llm_cond_dim=896,
	cfg_dropout_prob=0.1,
	**kwargs,
	):
	super().__init__()

	self.in_channels = in_channels
	self.out_channels = in_channels
	self.num_heads = num_heads


	self.t_embedder = TimestepEmbedder(hidden_size)
	self.x_embedder = nn.Linear(in_channels, hidden_size)
	self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size, cfg_dropout_prob)
	self.hidden_size = hidden_size

	self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)


	self.blocks = nn.ModuleList([
	DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)
	])
	self.final_layer = FinalLayer(hidden_size, self.out_channels)
	self.initialize_weights()


	def initialize_weights(self):
	# Initialize transformer layers:
	def _basic_init(module):
	if isinstance(module, nn.Linear):
	torch.nn.init.xavier_uniform_(module.weight)
	if module.bias is not None:
	nn.init.constant_(module.bias, 0)
	self.apply(_basic_init)


	# Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
	w_x = self.x_embedder.weight.data
	nn.init.xavier_uniform_(w_x.view([w_x.shape[0], -1]))
	nn.init.constant_(self.x_embedder.bias, 0)

	# Initialize label embedding table:
	w_c = self.c_embedder.cond_embedder.weight.data
	nn.init.xavier_uniform_(w_c.view([w_c.shape[0], -1]))
	nn.init.constant_(self.c_embedder.cond_embedder.bias, 0)

	# Initialize timestep embedding MLP:
	nn.init.normal_(self.t_embedder.time_mlp[0].weight, std=0.02)
	nn.init.normal_(self.t_embedder.time_mlp[2].weight, std=0.02)


	# Zero-out output layers:
	nn.init.constant_(self.final_layer.linear.weight, 0)
	nn.init.constant_(self.final_layer.linear.bias, 0)



	def forward(self, x, t, c, latent_history, mask=None):
	"""
	Forward pass of DiT.
	x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
	t: (N,) tensor of diffusion timesteps
	y: (N,) tensor of class labels
	"""
	t = self.t_embedder(t).unsqueeze(1) # (N, D)
	x_now = self.x_embedder(x) #转换x_noise的维度
	x_history = self.x_embedder(latent_history) #转换target history的维度
	x = torch.cat([x_history, x_now], dim=1) # 拼接patch history和当前patch
	c = self.c_embedder(c, self.training) # (N, 1, 896) -> (N, 1, D)，同时带有CFG替换操作
	y = t + c
	x = torch.cat([y, x], dim=1) # # (N, 1 + patch_size *2, D)
	rope = self.rotary_embed.forward_from_seq_len(x.shape[1])

	if mask is not None:
	mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1])
	mask = torch.cat([mask_pad, mask], dim=-1)
	for block in self.blocks:
	x = block(x, mask, rope) # (N, T, D)
	x = self.final_layer(x) # (N, T, patch_size ** 2 * out_channels)
	#x = self.unpatchify(x) # (N, out_channels, H, W)
	return x

	def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
	"""
	Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
	"""
	# https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
	if not cfg_scale == 1:
	x = torch.cat([x, x], dim=0)
	latent_history = torch.cat([latent_history, latent_history], dim=0)
	fake_latent = torch.zeros(c.shape).to(c.device)
	c = torch.cat([c, fake_latent], dim=0)
	if t.ndim == 0:
	t = t.repeat(x.shape[0])
	model_out = self.forward(x, t, c, latent_history)
	return model_out[:, -patch_size:, :]