Spaces:

VinayHajare
/

Text-To-Image-EfficientCLIP-GAN

Running

App Files Files Community

Text-To-Image-EfficientCLIP-GAN / model.py

VinayHajare

Upload 4 files

574a477 verified 3 months ago

raw history blame contribute delete

No virus

33.6 kB

	import torch
	import torch.nn as nn
	import numpy as np
	import torch.nn.functional as F
	from collections import OrderedDict
	from utils import dummy_context_mgr


	class CLIP_IMG_ENCODER(nn.Module):
	"""
	CLIP_IMG_ENCODER module for encoding images using CLIP's visual transformer.
	"""

	def __init__(self, CLIP):
	"""
	Initialize the CLIP_IMG_ENCODER module.

	Args:
	CLIP (CLIP): Pre-trained CLIP model.
	"""
	super(CLIP_IMG_ENCODER, self).__init__()
	model = CLIP.visual
	self.define_module(model)
	# freeze the parameters of the CLIP model
	for param in self.parameters():
	param.requires_grad = False

	def define_module(self, model):
	"""
	Define the individual layers and modules of the CLIP visual transformer model.
	Args:
	model (nn.Module): CLIP visual transformer model.
	"""
	# Extract required modules from the CLIP model
	self.conv1 = model.conv1 # Convolutional layer
	self.class_embedding = model.class_embedding # Class embedding layer
	self.positional_embedding = model.positional_embedding # Positional embedding layer
	self.ln_pre = model.ln_pre # Linear Normalization layer for pre-normalization
	self.transformer = model.transformer # Transformer block
	self.ln_post = model.ln_post # Linear Normalization layer for post-normalization
	self.proj = model.proj # projection matrix

	@property
	def dtype(self):
	"""
	Get the data type of the convolutional layer weights.
	"""
	return self.conv1.weight.dtype

	def transf_to_CLIP_input(self, inputs):
	"""
	Transform input images to the format expected by CLIP.

	Args:
	inputs (torch.Tensor): Input images.

	Returns:
	torch.Tensor: Transformed images.
	"""
	device = inputs.device
	# Check the size of the input image tensor
	if len(inputs.size()) != 4:
	raise ValueError('Expect the (B, C, X, Y) tensor.')
	else:
	# Normalize input images
	mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0).to(device)
	var = torch.tensor([0.26862954, 0.26130258, 0.27577711]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0).to(device)
	inputs = F.interpolate(inputs * 0.5 + 0.5, size=(224, 224))
	inputs = ((inputs + 1) * 0.5 - mean) / var
	return inputs

	def forward(self, img: torch.Tensor):
	"""
	Forward pass of the CLIP_IMG_ENCODER module.

	Args:
	img (torch.Tensor): Input images.

	Returns:
	torch.Tensor: Local features extracted from the image.
	torch.Tensor: Encoded image embeddings.
	"""
	# Transform input images to the format expected by CLIP and set its datatype appropriately
	x = self.transf_to_CLIP_input(img)
	x = x.type(self.dtype)

	# Pass the image through Convolutional layer
	x = self.conv1(x) # shape = [*, width, grid, grid]
	grid = x.size(-1)

	# Reshape and permute the tensor for transformer input
	x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [, width, grid * 2]
	x = x.permute(0, 2, 1) # shape = [, grid * 2, width]

	# Add class and positional embeddings
	x = torch.cat(
	[self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
	x], dim=1) # shape = [, grid * 2 + 1, width]
	x = x + self.positional_embedding.to(x.dtype)
	x = self.ln_pre(x)

	# NLD (Batch Size - Length - Dimension) -> LND (Length - Batch Size - Dimension)
	x = x.permute(1, 0, 2)

	# Extract local features using transformer blocks
	selected = [1, 4, 8]
	local_features = []
	for i in range(12):
	x = self.transformer.resblocks[i](x)
	if i in selected:
	local_features.append(
	x.permute(1, 0, 2)[:, 1:, :].permute(0, 2, 1).reshape(-1, 768, grid, grid).contiguous().type(
	img.dtype))
	x = x.permute(1, 0, 2) # LND -> NLD
	x = self.ln_post(x[:, 0, :])
	if self.proj is not None:
	x = x @ self.proj # Perform matrix multiplication with projection matrix and tensor
	return torch.stack(local_features, dim=1), x.type(img.dtype)


	class CLIP_TXT_ENCODER(nn.Module):
	"""
	CLIP_TXT_ENCODER module for encoding text inputs using CLIP's transformer.
	"""

	def __init__(self, CLIP):
	"""
	Initialize the CLIP_TXT_ENCODER module.

	Args:
	CLIP (CLIP): Pre-trained CLIP model.
	"""
	super(CLIP_TXT_ENCODER, self).__init__()
	self.define_module(CLIP)
	# Freeze the parameters of the CLIP model
	for param in self.parameters():
	param.requires_grad = False

	def define_module(self, CLIP):
	"""
	Define the individual modules of the CLIP transformer model.

	Args:
	CLIP (CLIP): Pre-trained CLIP model.
	"""
	self.transformer = CLIP.transformer # Transformer block
	self.vocab_size = CLIP.vocab_size # Size of the vocabulary of the transformer
	self.token_embedding = CLIP.token_embedding # token embedding block
	self.positional_embedding = CLIP.positional_embedding # positional embedding block
	self.ln_final = CLIP.ln_final # Linear Normalization layer
	self.text_projection = CLIP.text_projection # Projection matrix for text

	@property
	def dtype(self):
	"""
	Get the data type of the first layer's weights in the transformer.
	"""
	return self.transformer.resblocks[0].mlp.c_fc.weight.dtype

	def forward(self, text):
	"""
	Forward pass of the CLIP_TXT_ENCODER module.

	Args:
	text (torch.Tensor): Input text tokens.

	Returns:
	torch.Tensor: Encoded sentence embeddings.
	torch.Tensor: Transformer output for the input text.
	"""
	# Embed input text tokens
	x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model]
	# Add positional embeddings
	x = x + self.positional_embedding.type(self.dtype)
	# Permute dimensions for transformer input
	x = x.permute(1, 0, 2) # NLD -> LND
	# Pass input through the transformer
	x = self.transformer(x)
	# Permute dimensions back to original shape
	x = x.permute(1, 0, 2) # LND -> NLD
	# Apply layer normalization
	x = self.ln_final(x).type(self.dtype) # shape = [batch_size, n_ctx, transformer.width]
	# Extract sentence embeddings from the end-of-text (eot_token : is the highest number in each sequence)
	sent_emb = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection

	# Return the sentence embedding and transformer ouput
	return sent_emb, x


	class CLIP_Mapper(nn.Module):
	"""
	CLIP_Mapper module for mapping images with prompts using CLIP's transformer.
	"""

	def __init__(self, CLIP):
	"""
	Initialize the CLIP_Mapper module.

	Args:
	CLIP (CLIP): Pre-trained CLIP model.
	"""
	super(CLIP_Mapper, self).__init__()
	model = CLIP.visual
	self.define_module(model)
	# Freeze the parameters of the CLIP visual model
	for param in model.parameters():
	param.requires_grad = False

	def define_module(self, model):
	"""
	Define the individual modules of the CLIP visual model.

	Args:
	model: Pre-trained CLIP visual model.
	"""
	self.conv1 = model.conv1
	self.class_embedding = model.class_embedding
	self.positional_embedding = model.positional_embedding
	self.ln_pre = model.ln_pre
	self.transformer = model.transformer

	@property
	def dtype(self):
	"""
	Get the data type of the weights of the first convolutional layer.
	"""
	return self.conv1.weight.dtype

	def forward(self, img: torch.Tensor, prompts: torch.Tensor):
	"""
	Forward pass of the CLIP_Mapper module.

	Args:
	img (torch.Tensor): Input image tensor.
	prompts (torch.Tensor): Prompt tokens for mapping.

	Returns:
	torch.Tensor: Mapped features from the CLIP model.
	"""

	# Convert input image and prompts to the appropriate data type
	x = img.type(self.dtype)
	prompts = prompts.type(self.dtype)
	grid = x.size(-1)

	# Reshape the input image tensor
	x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [, width, grid * 2]
	x = x.permute(0, 2, 1) # shape = [, grid * 2, width]

	# Append the class embeddings to input tensors
	x = torch.cat(
	[self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
	x],
	dim=1
	) # shape = [, grid * 2 + 1, width]

	# Append the positional embeddings to the input tensor
	x = x + self.positional_embedding.to(x.dtype)

	# Perform the layer normalization
	x = self.ln_pre(x)
	# NLD -> LND
	x = x.permute(1, 0, 2)
	# Local features
	selected = [1, 2, 3, 4, 5, 6, 7, 8]
	begin, end = 0, 12
	prompt_idx = 0
	for i in range(begin, end):
	# Add prompt to the input tensor
	if i in selected:
	prompt = prompts[:, prompt_idx, :].unsqueeze(0)
	prompt_idx = prompt_idx + 1
	x = torch.cat((x, prompt), dim=0)
	x = self.transformer.resblocks[i](x)
	x = x[:-1, :, :]
	else:
	x = self.transformer.resblocks[i](x)
	# Reshape and return mapped features
	return x.permute(1, 0, 2)[:, 1:, :].permute(0, 2, 1).reshape(-1, 768, grid, grid).contiguous().type(img.dtype)


	class CLIP_Adapter(nn.Module):
	"""
	CLIP_Adapter module for adapting features from a generator to match the CLIP model's input requirements.
	"""

	def __init__(self, in_ch, mid_ch, out_ch, G_ch, CLIP_ch, cond_dim, k, s, p, map_num, CLIP):
	"""
	Initialize the CLIP_Adapter module.

	Args:
	in_ch (int): Number of input channels.
	mid_ch (int): Number of channels in the intermediate layers.
	out_ch (int): Number of output channels.
	G_ch (int): Number of channels in the generator's output.
	CLIP_ch (int): Number of channels in the CLIP model's input.
	cond_dim (int): Dimension of the conditioning vector.
	k (int): Kernel size for convolutional layers.
	s (int): Stride for convolutional layers.
	p (int): Padding for convolutional layers.
	map_num (int): Number of mapping blocks.
	CLIP: Pre-trained CLIP model.
	"""
	super(CLIP_Adapter, self).__init__()
	self.CLIP_ch = CLIP_ch
	self.FBlocks = nn.ModuleList([])
	# Define Mapping blocks (M_Block) and them to Feature blocks (FBlock) for given number of mapping blocks.
	self.FBlocks.append(M_Block(in_ch, mid_ch, out_ch, cond_dim, k, s, p))
	for i in range(map_num - 1):
	self.FBlocks.append(M_Block(out_ch, mid_ch, out_ch, cond_dim, k, s, p))
	# Convolutional layer to fuse adapted features
	self.conv_fuse = nn.Conv2d(out_ch, CLIP_ch, 5, 1, 2)
	# CLIP Mapper module to map adapted features to CLIP's input space
	self.CLIP_ViT = CLIP_Mapper(CLIP)
	# Convolutional layer to further process mapped features
	self.conv = nn.Conv2d(768, G_ch, 5, 1, 2)
	# Fully connected layer for conditioning
	self.fc_prompt = nn.Linear(cond_dim, CLIP_ch * 8)

	def forward(self, out, c):
	"""
	Forward pass of the CLIP_Adapter module. Takes output features from the generator and conditioning vector
	as input, adapts features using the Feature block having multiple mapping blocks, fuses them, map them to
	CLIPs input space and returns the processed features

	Args:
	out (torch.Tensor): Output features from the generator.
	c (torch.Tensor): Conditioning vector.

	Returns:
	torch.Tensor: Adapted and mapped features for the generator.
	"""

	# Generate prompts from the conditioning vector
	prompts = self.fc_prompt(c).view(c.size(0), -1, self.CLIP_ch)

	# Pass features through feature block consisting of multiple mapping blocks
	for FBlock in self.FBlocks:
	out = FBlock(out, c)
	# Fuse adapted features
	fuse_feat = self.conv_fuse(out)
	# Map fused features to CLIP's input space
	map_feat = self.CLIP_ViT(fuse_feat, prompts)
	# Further process mapped features and return
	return self.conv(fuse_feat + 0.1 * map_feat)


	class NetG(nn.Module):
	"""
	Generator network for synthesizing images conditioned on text and noise
	"""

	def __init__(self, ngf, nz, cond_dim, imsize, ch_size, mixed_precision, CLIP):
	"""
	Initializes the Generator network.

	Parameters:
	ngf (int): Number of generator filters.
	nz (int): Dimensionality of the input noise vector.
	cond_dim (int): Dimensionality of the conditioning vector.
	imsize (int): Size of the generated images.
	ch_size (int): Number of output channels for the generated images.
	mixed_precision (bool): Whether to use mixed precision training.
	CLIP: CLIP model for feature adaptation.

	"""
	super(NetG, self).__init__()
	# Define attributes
	self.ngf = ngf
	self.mixed_precision = mixed_precision

	# Build CLIP Mapper
	self.code_sz, self.code_ch, self.mid_ch = 7, 64, 32
	self.CLIP_ch = 768
	# fully connected layer to convert the noise vector into a feature map of dimensions (code_sz * code_sz * code_ch)
	self.fc_code = nn.Linear(nz, self.code_sz * self.code_sz * self.code_ch)
	self.mapping = CLIP_Adapter(self.code_ch, self.mid_ch, self.code_ch, ngf * 8, self.CLIP_ch, cond_dim + nz, 3, 1,
	1, 4, CLIP)
	# Build GBlocks
	self.GBlocks = nn.ModuleList([])
	in_out_pairs = list(get_G_in_out_chs(ngf, imsize))
	imsize = 4
	for idx, (in_ch, out_ch) in enumerate(in_out_pairs):
	if idx < (len(in_out_pairs) - 1):
	imsize = imsize * 2
	else:
	imsize = 224
	self.GBlocks.append(G_Block(cond_dim + nz, in_ch, out_ch, imsize))

	# To RGB image conversion using the sequential layers having leakyReLU activation function
	self.to_rgb = nn.Sequential(
	nn.LeakyReLU(0.2, inplace=True),
	nn.Conv2d(out_ch, ch_size, 3, 1, 1),
	)

	def forward(self, noise, c, eval=False): # x=noise, c=ent_emb
	"""
	Forward pass of the generator network.

	Args:
	noise (torch.Tensor): Input noise vector.
	c (torch.Tensor): Conditioning information, typically an embedding representing attributes of the output.
	eval (bool, optional): Flag indicating whether the network is in evaluation mode. Defaults to False.

	Returns:
	torch.Tensor: Generated RGB images.
	"""
	# Context manager for enabling automatic mixed precision training
	with torch.cuda.amp.autocast() if self.mixed_precision and not eval else dummy_context_mgr() as mp:
	# Concatenate noise and conditioning information
	cond = torch.cat((noise, c), dim=1)

	# Pass noise through fully connected layer to generate feature map and adapt features using CLIP Mapper
	out = self.mapping(self.fc_code(noise).view(noise.size(0), self.code_ch, self.code_sz, self.code_sz), cond)

	# Apply GBlocks to progressively upsample feature representation, fuse text and visual features
	for GBlock in self.GBlocks:
	out = GBlock(out, cond)

	# Convert final feature representation to RGB images
	out = self.to_rgb(out)

	return out


	class NetD(nn.Module):
	"""
	Discriminator network for evaluating the realism of images.
	Attributes:
	DBlocks (nn.ModuleList): List of D_Block modules for processing feature maps.
	main (D_Block): Main D_Block module for final processing.
	"""

	def __init__(self, ndf, imsize, ch_size, mixed_precision):
	"""
	Initializes the Discriminator network

	Args:
	ndf (int): Number of channels in the initial features.
	imsize (int): Size of the input images (assumed square).
	ch_size (int): Number of channels in the output feature maps.
	mixed_precision (bool): Flag indicating whether to use mixed precision training.
	"""
	super(NetD, self).__init__()
	self.mixed_precision = mixed_precision
	# Define the DBlock
	self.DBlocks = nn.ModuleList([
	D_Block(768, 768, 3, 1, 1, res=True, CLIP_feat=True),
	D_Block(768, 768, 3, 1, 1, res=True, CLIP_feat=True),
	])
	# Define the main DBlock for the final processing
	self.main = D_Block(768, 512, 3, 1, 1, res=True, CLIP_feat=False)

	def forward(self, h):
	"""
	Forward pass of the discriminator network.
	Args:
	h (torch.Tensor): Input feature maps.
	Returns:
	torch.Tensor: Discriminator output.
	"""
	with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
	# Initial feature map
	out = h[:, 0]
	# Pass the input feature through each DBlock
	for idx in range(len(self.DBlocks)):
	out = self.DBlocks[idx](out, h[:, idx + 1])
	# Final processing through the main DBlock
	out = self.main(out)
	return out


	class NetC(nn.Module):
	"""
	Classifier / Comparator network for classifying the joint features of the generator output and condition text.
	Attributes:
	cond_dim (int): Dimensionality of the conditioning information.
	mixed_precision (bool): Flag indicating whether to use mixed precision training.
	joint_conv (nn.Sequential): Sequential module defining the classifier layers.
	"""
	def __init__(self, ndf, cond_dim, mixed_precision):
	"""

	"""
	super(NetC, self).__init__()
	self.cond_dim = cond_dim
	self.mixed_precision = mixed_precision
	# Define the classifier layers, sequential convolutional 2D layer with LeakyReLU as the activation function
	self.joint_conv = nn.Sequential(
	nn.Conv2d(512 + 512, 128, 4, 1, 0, bias=False),
	nn.LeakyReLU(0.2, inplace=True),
	nn.Conv2d(128, 1, 4, 1, 0, bias=False),
	)

	def forward(self, out, cond):
	"""
	Forward pass of the classifier network.

	Args:
	out (torch.Tensor): Generator output feature map.
	cond (torch.Tensor): Conditioning information vector
	"""
	with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
	# Reshape and repeat conditioning information vector to match the feature map size
	cond = cond.view(-1, self.cond_dim, 1, 1)
	cond = cond.repeat(1, 1, 7, 7)

	# Concatenate feature map and conditioned information
	h_c_code = torch.cat((out, cond), 1)

	# Pass through the classifier layers
	out = self.joint_conv(h_c_code)
	return out


	class M_Block(nn.Module):
	"""
	Multi-scale block consisting of convolutional layers and conditioning.

	Attributes:
	conv1 (nn.Conv2d): First convolutional layer.
	fuse1 (DFBlock): Conditioning block for the first convolutional layer.
	conv2 (nn.Conv2d): Second convolutional layer.
	fuse2 (DFBlock): Conditioning block for the second convolutional layer.
	learnable_sc (bool): Flag indicating whether the shortcut connection is learnable.
	c_sc (nn.Conv2d): Convolutional layer for the shortcut connection.

	"""
	def __init__(self, in_ch, mid_ch, out_ch, cond_dim, k, s, p):
	"""
	Initializes the Multi-scale block.

	Args:
	in_ch (int): Number of input channels.
	mid_ch (int): Number of channels in the intermediate layers.
	out_ch (int): Number of output channels.
	cond_dim (int): Dimensionality of the conditioning information.
	k (int): Kernel size for convolutional layers.
	s (int): Stride for convolutional layers.
	p (int): Padding for convolutional layers.

	"""
	super(M_Block, self).__init__()

	# Define convolutional layers and conditioning blocks
	self.conv1 = nn.Conv2d(in_ch, mid_ch, k, s, p)
	self.fuse1 = DFBLK(cond_dim, mid_ch)
	self.conv2 = nn.Conv2d(mid_ch, out_ch, k, s, p)
	self.fuse2 = DFBLK(cond_dim, out_ch)

	# Learnable shortcut connection
	self.learnable_sc = in_ch != out_ch
	if self.learnable_sc:
	self.c_sc = nn.Conv2d(in_ch, out_ch, 1, stride=1, padding=0)

	def shortcut(self, x):
	"""
	Defines the shortcut connection.

	Args:
	x (torch.Tensor): Input tensor.

	Returns:
	torch.Tensor: Shortcut connection output.
	"""
	if self.learnable_sc:
	x = self.c_sc(x)
	return x

	def residual(self, h, text):
	"""
	Defines the residual path with conditioning.

	Args:
	h (torch.Tensor): Input tensor.
	text (torch.Tensor): Conditioning information.

	Returns:
	torch.Tensor: Residual path output.
	"""
	h = self.conv1(h)
	h = self.fuse1(h, text)
	h = self.conv2(h)
	h = self.fuse2(h, text)
	return h

	def forward(self, h, c):
	"""
	Forward pass of the multi-scale block.

	Args:
	h (torch.Tensor): Input tensor.
	c (torch.Tensor): Conditioning information.

	Returns:
	torch.Tensor: Output tensor.
	"""
	return self.shortcut(h) + self.residual(h, c)


	class G_Block(nn.Module):
	"""
	Generator block consisting of convolutional layers and conditioning.

	Attributes:
	imsize (int): Size of the output image.
	learnable_sc (bool): Flag indicating whether the shortcut connection is learnable.
	c1 (nn.Conv2d): First convolutional layer.
	c2 (nn.Conv2d): Second convolutional layer.
	fuse1 (DFBLK): Conditioning block for the first convolutional layer.
	fuse2 (DFBLK): Conditioning block for the second convolutional layer.
	c_sc (nn.Conv2d): Convolutional layer for the shortcut connection.
	"""

	def __init__(self, cond_dim, in_ch, out_ch, imsize):
	"""
	Initialize the Generator block.

	Args:
	cond_dim (int): Dimensionality of the conditioning information.
	in_ch (int): Number of input channels.
	out_ch (int): Number of output channels.
	imsize (int): Size of the output image.
	"""
	super(G_Block, self).__init__()

	# Initialize attributes
	self.imsize = imsize
	self.learnable_sc = in_ch != out_ch

	# Define convolution layers and conditioning blocks
	self.c1 = nn.Conv2d(in_ch, out_ch, 3, 1, 1)
	self.c2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1)
	self.fuse1 = DFBLK(cond_dim, in_ch)
	self.fuse2 = DFBLK(cond_dim, out_ch)

	# Learnable shortcut connection
	if self.learnable_sc:
	self.c_sc = nn.Conv2d(in_ch, out_ch, 1, stride=1, padding=0)

	def shortcut(self, x):
	"""
	Defines the shortcut connection.

	Args:
	x (torch.Tensor): Input tensor.

	Returns:
	torch.Tensor: Shortcut connection output.
	"""
	if self.learnable_sc:
	x = self.c_sc(x)
	return x

	def residual(self, h, y):
	"""
	Defines the residual path with conditioning.

	Args:
	h (torch.Tensor): Input tensor.
	y (torch.Tensor): Conditioning information.

	Returns:
	torch.Tensor: Residual path output.
	"""
	h = self.fuse1(h, y)
	h = self.c1(h)
	h = self.fuse2(h, y)
	h = self.c2(h)
	return h

	def forward(self, h, y):
	"""
	Forward pass of the generator block.

	Args:
	h (torch.Tensor): Input tensor.
	y (torch.Tensor): Conditioning information.

	Returns:
	torch.Tensor: Output tensor.
	"""
	h = F.interpolate(h, size=(self.imsize, self.imsize))
	return self.shortcut(h) + self.residual(h, y)


	class D_Block(nn.Module):
	"""
	Discriminator block.
	"""
	def __init__(self, fin, fout, k, s, p, res, CLIP_feat):
	"""
	Initializes Discriminator block.

	Args:
	- fin (int): Number of input channels.
	- fout (int): Number of output channels.
	- k (int): Kernel size for convolutional layers.
	- s (int): Stride for convolutional layers.
	- p (int): Padding for convolutional layers.
	- res (bool): Whether to use residual connection.
	- CLIP_feat (bool): Whether to incorporate CLIP features.
	"""
	super(D_Block, self).__init__()
	self.res, self.CLIP_feat = res, CLIP_feat
	self.learned_shortcut = (fin != fout)

	# Convolutional layers for residual path
	self.conv_r = nn.Sequential(
	nn.Conv2d(fin, fout, k, s, p, bias=False),
	nn.LeakyReLU(0.2, inplace=True),
	nn.Conv2d(fout, fout, k, s, p, bias=False),
	nn.LeakyReLU(0.2, inplace=True),
	)

	# Convolutional layers for shortcut connection
	self.conv_s = nn.Conv2d(fin, fout, 1, stride=1, padding=0)

	# Parameters for learned residual and CLIP features
	if self.res == True:
	self.gamma = nn.Parameter(torch.zeros(1))
	if self.CLIP_feat == True:
	self.beta = nn.Parameter(torch.zeros(1))

	def forward(self, x, CLIP_feat=None):
	"""
	Forward pass of the discriminator block.

	Args:
	- x (torch.Tensor): Input tensor.
	- CLIP_feat (torch.Tensor): Optional CLIP features tensor.

	Returns:
	- torch.Tensor: Output tensor.
	"""
	# Compute the residual features
	res = self.conv_r(x)

	# Compute the shortcut connection
	if self.learned_shortcut:
	x = self.conv_s(x)

	# Incorporate learned residual and CLIP features if enabled
	if (self.res == True) and (self.CLIP_feat == True):
	return x + self.gamma * res + self.beta * CLIP_feat
	elif (self.res == True) and (self.CLIP_feat != True):
	return x + self.gamma * res
	elif (self.res != True) and (self.CLIP_feat == True):
	return x + self.beta * CLIP_feat
	else:
	return x


	class DFBLK(nn.Module):
	"""
	Diffusion Block of the Generator network with Conditional feature block
	"""
	def __init__(self, cond_dim, in_ch):
	"""
	Initializing the Conditional feature block of the DFBlock.

	Args:
	- cond_dim (int): Dimensionality of the conditional input.
	- in_ch (int): Number of input channels.
	"""
	super(DFBLK, self).__init__()
	# Define conditional affine transformations
	self.affine0 = Affine(cond_dim, in_ch)
	self.affine1 = Affine(cond_dim, in_ch)

	def forward(self, x, y=None):
	"""
	Forward pass of the conditional feature block.

	Args:
	- x (torch.Tensor): Input tensor.
	- y (torch.Tensor, optional): Conditional input tensor. Default is None.

	Returns:
	- torch.Tensor: Output tensor.
	"""
	# Apply the first affine transformation and activation function
	h = self.affine0(x, y)
	h = nn.LeakyReLU(0.2, inplace=True)(h)
	# Apply second affine transformation and activation function
	h = self.affine1(h, y)
	h = nn.LeakyReLU(0.2, inplace=True)(h)
	return h


	class QuickGELU(nn.Module):
	"""
	Efficient and faster version of GELU,
	for non-linearity and to learn complex patterns
	"""
	def forward(self, x: torch.Tensor):
	"""
	Forward pass of the QuickGELU activation function.

	Args:
	- x (torch.Tensor): Input tensor.

	Returns:
	- torch.Tensor: Output tensor.
	"""
	# Apply QuickGELU activation function
	return x * torch.sigmoid(1.702 * x)


	# Taken from the RAT-GAN repository
	class Affine(nn.Module):
	"""
	Affine transformation module that applies conditional scaling and shifting to input features,
	to incorporate additional control over the generated output based on input conditions.
	"""
	def __init__(self, cond_dim, num_features):
	"""
	Initialize the affine transformation module.
	Args:
	cond_dim (int): Dimensionality of the conditioning information.
	num_features (int): Number of input features.
	"""
	super(Affine, self).__init__()
	# Define 2 fully connected networks to compute gamma and beta parameters
	# each 2 linear layers with RELU activation in between
	self.fc_gamma = nn.Sequential(OrderedDict([
	('linear1', nn.Linear(cond_dim, num_features)),
	('relu1', nn.ReLU(inplace=True)),
	('linear2', nn.Linear(num_features, num_features)),
	]))
	self.fc_beta = nn.Sequential(OrderedDict([
	('linear1', nn.Linear(cond_dim, num_features)),
	('relu1', nn.ReLU(inplace=True)),
	('linear2', nn.Linear(num_features, num_features)),
	]))
	# Initializes the weights and biases of the network
	self._initialize()

	def _initialize(self):
	"""
	Initializes the weights and biases of the linear layers responsible for computing gamma and beta
	"""
	nn.init.zeros_(self.fc_gamma.linear2.weight.data)
	nn.init.ones_(self.fc_gamma.linear2.bias.data)
	nn.init.zeros_(self.fc_beta.linear2.weight.data)
	nn.init.zeros_(self.fc_beta.linear2.bias.data)

	def forward(self, x, y=None):
	"""
	Forward pass of the Affine transformation module.

	Args:
	x (torch.Tensor): Input tensor.
	y (torch.Tensor, optional): Conditioning information tensor. Default is None.

	Returns:
	torch.Tensor: Transformed tensor after applying affine transformation.
	"""
	# Compute gamma and beta parameters
	weight = self.fc_gamma(y)
	bias = self.fc_beta(y)

	# Ensure proper shape for weight and bias tensors
	if weight.dim() == 1:
	weight = weight.unsqueeze(0)
	if bias.dim() == 1:
	bias = bias.unsqueeze(0)

	# Expand weight and bias tensors to match input tensor shape
	size = x.size()
	weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
	bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)

	# Apply affine transformation
	return weight * x + bias


	def get_G_in_out_chs(nf, imsize):
	"""
	Compute input-output channel pairs for generator blocks based on given number of channels and image size.

	Args:
	nf (int): Number of input channels.
	imsize (int): Size of the input image.

	Returns:
	list: List of tuples containing input-output channel pairs for generator blocks.
	"""
	# Determine the number of layers based on image size
	layer_num = int(np.log2(imsize)) - 1

	# Compute the number of channels for each layer
	channel_nums = [nf * min(2 ** idx, 8) for idx in range(layer_num)]

	# Reverse the channel numbers to start with the highest channel count
	channel_nums = channel_nums[::-1]

	# Generate input-output channel pairs for generator blocks
	in_out_pairs = zip(channel_nums[:-1], channel_nums[1:])

	return in_out_pairs


	def get_D_in_out_chs(nf, imsize):
	"""
	Compute input-output channel pairs for discriminator blocks based on given number of channels and image size.

	Args:
	nf (int): Number of input channels.
	imsize (int): Size of the input image.

	Returns:
	list: List of tuples containing input-output channel pairs for discriminator blocks.
	"""
	# Determine the number of layers based on image size
	layer_num = int(np.log2(imsize)) - 1

	# Compute the number of channels for each layer
	channel_nums = [nf * min(2 ** idx, 8) for idx in range(layer_num)]

	# Generate input-output channel pairs for discriminator blocks
	in_out_pairs = zip(channel_nums[:-1], channel_nums[1:])

	return in_out_pairs