TAI Research

Initial commit: Lumina_Dev_Legacy (archived)

29691f6 5 months ago

14.6 kB

import torch
import torch.nn as nn
from typing import Optional, Tuple, List, Union
import numpy as np
from tqdm import tqdm
import math


class DDIMSampler:
    """DDIM采样器"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 50):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
        
        # 设置时间步
        self.set_timesteps(num_inference_steps)
    
    def set_timesteps(self, num_inference_steps: int):
        """设置推理时间步"""
        self.num_inference_steps = num_inference_steps
        
        # 选择时间步
        if self.diffusion.num_train_timesteps == num_inference_steps:
            timesteps = np.arange(0, self.diffusion.num_train_timesteps)
        else:
            step_ratio = self.diffusion.num_train_timesteps // num_inference_steps
            timesteps = np.arange(0, self.diffusion.num_train_timesteps, step_ratio)
        
        self.timesteps = torch.from_numpy(timesteps).long().flip(0)
    
    @torch.no_grad()
    def step(
        self,
        model_output: torch.Tensor,
        timestep: int,
        sample: torch.Tensor,
        eta: float = 0.0,
        use_clipped_model_output: bool = False
    ) -> torch.Tensor:
        """DDIM单步采样"""
        # 获取当前和上一个时间步
        prev_timestep = timestep - self.diffusion.num_train_timesteps // self.num_inference_steps
        
        # 提取alpha参数
        alpha_prod_t = self.diffusion.extract(self.diffusion.alphas_cumprod, timestep, sample.shape)
        alpha_prod_t_prev = self.diffusion.extract(
            self.diffusion.alphas_cumprod, 
            prev_timestep, 
            sample.shape
        ) if prev_timestep >= 0 else torch.ones_like(alpha_prod_t)
        
        # 根据预测类型处理模型输出
        if self.diffusion.prediction_type == "epsilon":
            pred_original_sample = (sample - (1 - alpha_prod_t) ** 0.5 * model_output) / alpha_prod_t ** 0.5
            pred_epsilon = model_output
        elif self.diffusion.prediction_type == "sample":
            pred_original_sample = model_output
            pred_epsilon = (sample - alpha_prod_t ** 0.5 * pred_original_sample) / (1 - alpha_prod_t) ** 0.5
        elif self.diffusion.prediction_type == "v_prediction":
            pred_original_sample = (alpha_prod_t ** 0.5) * sample - (1 - alpha_prod_t) ** 0.5 * model_output
            pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (1 - alpha_prod_t) ** 0.5 * sample
        else:
            raise ValueError(f"Unsupported prediction type: {self.diffusion.prediction_type}")
        
        # 裁剪预测的原始样本
        if use_clipped_model_output:
            pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
        
        # 计算x_t-1的方差
        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
        std_dev_t = eta * variance ** 0.5
        
        # 当eta > 0时，使用随机采样
        if eta > 0:
            noise = torch.randn_like(model_output)
            variance = std_dev_t ** 2
        else:
            noise = 0
            variance = 0
        
        # 计算x_t-1的均值
        pred_sample_direction = (1 - alpha_prod_t_prev - variance) ** 0.5 * pred_epsilon
        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
        
        # 添加噪声
        if eta > 0:
            prev_sample = prev_sample + std_dev_t * noise
        
        return prev_sample
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        height: int = 512,
        width: int = 512,
        num_images_per_prompt: int = 1,
        guidance_scale: float = 7.5,
        eta: float = 0.0,
        generator: Optional[torch.Generator] = None,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """生成样本"""
        # 设置模型为评估模式
        self.model.eval()
        
        # 批次大小
        batch_size = prompt_embeds.shape[0]
        
        # 初始化潜在表示
        latents = torch.randn(
            (batch_size * num_images_per_prompt, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device,
            generator=generator
        )
        
        # 准备额外的条件
        if negative_prompt_embeds is not None:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        
        # 分类器自由引导的缩放因子
        do_classifier_free_guidance = guidance_scale > 1.0
        if do_classifier_free_guidance:
            latents = torch.cat([latents] * 2)
        
        # 采样循环
        timesteps = self.timesteps.to(latents.device)
        
        if progress_bar:
            timesteps = tqdm(timesteps, desc="DDIM Sampling")
        
        for i, t in enumerate(timesteps):
            # 扩展潜在表示以匹配引导的批次大小
            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
            latent_model_input = self.diffusion.scale_model_input(latent_model_input, t)
            
            # 预测噪声
            noise_pred = self.model(latent_model_input, t, prompt_embeds)
            
            # 执行分类器自由引导
            if do_classifier_free_guidance:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            
            # 计算上一个样本
            latents = self.step(noise_pred, t, latents, eta)
        
        return latents


class DPMSampler:
    """DPM采样器（更快）"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 20):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        height: int = 512,
        width: int = 512,
        guidance_scale: float = 7.5,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """DPM采样"""
        self.model.eval()
        
        # 初始化潜在表示
        latents = torch.randn(
            (1, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device
        )
        
        # 简化的DPM采样
        timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)
        
        if progress_bar:
            timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="DPM Sampling")
        else:
            timesteps_iter = enumerate(timesteps[:-1])
        
        for i, t in timesteps_iter:
            # 预测噪声
            noise_pred = self.model(latents, t.unsqueeze(0) * 999, prompt_embeds)
            
            # 应用引导
            if guidance_scale > 1.0:
                # 简单引导
                noise_pred = noise_pred * guidance_scale
            
            # DPM更新步骤
            dt = timesteps[i + 1] - t
            latents = latents + dt * noise_pred
        
        return latents


class LCMSampler:
    """LCM（潜在一致性模型）采样器，极快"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 4):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
        
        # LCM特定的参数
        self.c_skip = 1.0
        self.c_out = 1.0
        self.c_in = 1.0
        self.c_noise = 1.0
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        height: int = 512,
        width: int = 512,
        guidance_scale: float = 7.5,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """LCM采样（极快，只需要4-8步）"""
        self.model.eval()
        
        # 初始化潜在表示
        latents = torch.randn(
            (1, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device
        )
        
        # LCM采样循环
        timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)
        
        if progress_bar:
            timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="LCM Sampling")
        else:
            timesteps_iter = enumerate(timesteps[:-1])
        
        for i, t in timesteps_iter:
            # LCM特定的缩放
            c_skip = self.c_skip
            c_out = self.c_out
            c_in = self.c_in
            c_noise = self.c_noise
            
            # 缩放输入
            scaled_latents = c_in * latents
            
            # 预测
            noise_pred = self.model(scaled_latents, c_noise * t.unsqueeze(0), prompt_embeds)
            
            # LCM更新规则
            denoised = c_skip * latents + c_out * noise_pred
            
            # 更新潜在表示
            dt = timesteps[i + 1] - t
            latents = denoised + dt * noise_pred
        
        return latents


class SamplerFactory:
    """采样器工厂"""
    @staticmethod
    def create_sampler(
        sampler_type: str,
        model: nn.Module,
        diffusion,
        num_inference_steps: int = 50
    ):
        """创建采样器"""
        if sampler_type == "ddim":
            return DDIMSampler(model, diffusion, num_inference_steps)
        elif sampler_type == "dpm":
            return DPMSampler(model, diffusion, num_inference_steps)
        elif sampler_type == "lcm":
            return LCMSampler(model, diffusion, num_inference_steps)
        else:
            raise ValueError(f"未知的采样器类型: {sampler_type}")


class TextToImagePipeline:
    """文本到图像管道"""
    def __init__(
        self,
        model: nn.Module,
        diffusion,
        text_encoder,
        vae_decoder,
        sampler_type: str = "ddim",
        device: str = "cuda"
    ):
        self.model = model.to(device)
        self.diffusion = diffusion
        self.text_encoder = text_encoder
        self.vae_decoder = vae_decoder
        self.sampler_type = sampler_type
        self.device = device
        
        # 创建采样器
        self.sampler = SamplerFactory.create_sampler(
            sampler_type, model, diffusion
        )
        
        # 设置为评估模式
        self.model.eval()
        if self.vae_decoder is not None:
            self.vae_decoder.eval()
    
    @torch.no_grad()
    def __call__(
        self,
        prompt: str,
        negative_prompt: str = "",
        height: int = 512,
        width: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        num_images: int = 1,
        seed: Optional[int] = None,
        progress_bar: bool = True
    ) -> List:
        """生成图像"""
        # 设置随机种子
        if seed is not None:
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(seed)
        
        # 编码提示
        prompt_embeds = self.text_encoder.encode([prompt]).to(self.device)
        negative_prompt_embeds = None
        
        if negative_prompt:
            negative_prompt_embeds = self.text_encoder.encode([negative_prompt]).to(self.device)
        
        # 生成潜在表示
        latents = self.sampler.sample(
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            height=height,
            width=width,
            num_images_per_prompt=num_images,
            guidance_scale=guidance_scale,
            progress_bar=progress_bar
        )
        
        # 解码为图像
        images = []
        for i in range(num_images):
            latent = latents[i:i+1]
            
            if self.vae_decoder is not None:
                image = self.vae_decoder(latent)
            else:
                # 如果没有VAE解码器，返回潜在表示
                image = latent
            
            images.append(image.cpu())
        
        return images
    
    def generate_grid(
        self,
        prompts: List[str],
        grid_size: Tuple[int, int] = (2, 2),
        **kwargs
    ) -> torch.Tensor:
        """生成图像网格"""
        images = []
        
        for prompt in prompts[:grid_size[0] * grid_size[1]]:
            image = self(prompt, **kwargs)[0]
            images.append(image)
        
        # 创建网格
        from torchvision.utils import make_grid
        grid = make_grid(torch.cat(images, dim=0), nrow=grid_size[1])
        
        return grid


def test_sampler():
    """测试采样器"""
    import torch.nn as nn
    
    # 创建模拟模型
    class MockModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.in_channels = 4
        
        def forward(self, x, t, context):
            # 返回随机噪声
            return torch.randn_like(x)
    
    # 创建模拟扩散过程
    class MockDiffusion:
        def __init__(self):
            self.num_train_timesteps = 1000
            self.alphas_cumprod = torch.ones(1000)
            self.prediction_type = "epsilon"
        
        def extract(self, a, t, x_shape):
            return torch.ones(x_shape[0], 1, 1, 1)
        
        def scale_model_input(self, x, t):
            return x
    
    model = MockModel()
    diffusion = MockDiffusion()
    
    # 测试DDIM采样器
    sampler = DDIMSampler(model, diffusion, num_inference_steps=10)
    
    # 测试采样
    prompt_embeds = torch.randn(1, 77, 768)
    latents = sampler.sample(prompt_embeds, height=64, width=64, progress_bar=False)
    
    print(f"DDIM采样完成，潜在表示形状: {latents.shape}")
    
    return sampler, latents


if __name__ == '__main__':
    test_sampler()

	import torch
	import torch.nn as nn
	from typing import Optional, Tuple, List, Union
	import numpy as np
	from tqdm import tqdm
	import math


	class DDIMSampler:
	"""DDIM采样器"""
	def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 50):
	self.model = model
	self.diffusion = diffusion
	self.num_inference_steps = num_inference_steps

	# 设置时间步
	self.set_timesteps(num_inference_steps)

	def set_timesteps(self, num_inference_steps: int):
	"""设置推理时间步"""
	self.num_inference_steps = num_inference_steps

	# 选择时间步
	if self.diffusion.num_train_timesteps == num_inference_steps:
	timesteps = np.arange(0, self.diffusion.num_train_timesteps)
	else:
	step_ratio = self.diffusion.num_train_timesteps // num_inference_steps
	timesteps = np.arange(0, self.diffusion.num_train_timesteps, step_ratio)

	self.timesteps = torch.from_numpy(timesteps).long().flip(0)

	@torch.no_grad()
	def step(
	self,
	model_output: torch.Tensor,
	timestep: int,
	sample: torch.Tensor,
	eta: float = 0.0,
	use_clipped_model_output: bool = False
	) -> torch.Tensor:
	"""DDIM单步采样"""
	# 获取当前和上一个时间步
	prev_timestep = timestep - self.diffusion.num_train_timesteps // self.num_inference_steps

	# 提取alpha参数
	alpha_prod_t = self.diffusion.extract(self.diffusion.alphas_cumprod, timestep, sample.shape)
	alpha_prod_t_prev = self.diffusion.extract(
	self.diffusion.alphas_cumprod,
	prev_timestep,
	sample.shape
	) if prev_timestep >= 0 else torch.ones_like(alpha_prod_t)

	# 根据预测类型处理模型输出
	if self.diffusion.prediction_type == "epsilon":
	pred_original_sample = (sample - (1 - alpha_prod_t) ** 0.5 * model_output) / alpha_prod_t ** 0.5
	pred_epsilon = model_output
	elif self.diffusion.prediction_type == "sample":
	pred_original_sample = model_output
	pred_epsilon = (sample - alpha_prod_t ** 0.5 * pred_original_sample) / (1 - alpha_prod_t) ** 0.5
	elif self.diffusion.prediction_type == "v_prediction":
	pred_original_sample = (alpha_prod_t ** 0.5) * sample - (1 - alpha_prod_t) ** 0.5 * model_output
	pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (1 - alpha_prod_t) ** 0.5 * sample
	else:
	raise ValueError(f"Unsupported prediction type: {self.diffusion.prediction_type}")

	# 裁剪预测的原始样本
	if use_clipped_model_output:
	pred_original_sample = torch.clamp(pred_original_sample, -1, 1)

	# 计算x_t-1的方差
	variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
	std_dev_t = eta * variance ** 0.5

	# 当eta > 0时，使用随机采样
	if eta > 0:
	noise = torch.randn_like(model_output)
	variance = std_dev_t ** 2
	else:
	noise = 0
	variance = 0

	# 计算x_t-1的均值
	pred_sample_direction = (1 - alpha_prod_t_prev - variance) ** 0.5 * pred_epsilon
	prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

	# 添加噪声
	if eta > 0:
	prev_sample = prev_sample + std_dev_t * noise

	return prev_sample

	@torch.no_grad()
	def sample(
	self,
	prompt_embeds: torch.Tensor,
	negative_prompt_embeds: Optional[torch.Tensor] = None,
	height: int = 512,
	width: int = 512,
	num_images_per_prompt: int = 1,
	guidance_scale: float = 7.5,
	eta: float = 0.0,
	generator: Optional[torch.Generator] = None,
	progress_bar: bool = True
	) -> torch.Tensor:
	"""生成样本"""
	# 设置模型为评估模式
	self.model.eval()

	# 批次大小
	batch_size = prompt_embeds.shape[0]

	# 初始化潜在表示
	latents = torch.randn(
	(batch_size * num_images_per_prompt, self.model.in_channels, height // 8, width // 8),
	device=prompt_embeds.device,
	generator=generator
	)

	# 准备额外的条件
	if negative_prompt_embeds is not None:
	prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

	# 分类器自由引导的缩放因子
	do_classifier_free_guidance = guidance_scale > 1.0
	if do_classifier_free_guidance:
	latents = torch.cat([latents] * 2)

	# 采样循环
	timesteps = self.timesteps.to(latents.device)

	if progress_bar:
	timesteps = tqdm(timesteps, desc="DDIM Sampling")

	for i, t in enumerate(timesteps):
	# 扩展潜在表示以匹配引导的批次大小
	latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
	latent_model_input = self.diffusion.scale_model_input(latent_model_input, t)

	# 预测噪声
	noise_pred = self.model(latent_model_input, t, prompt_embeds)

	# 执行分类器自由引导
	if do_classifier_free_guidance:
	noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
	noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

	# 计算上一个样本
	latents = self.step(noise_pred, t, latents, eta)

	return latents


	class DPMSampler:
	"""DPM采样器（更快）"""
	def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 20):
	self.model = model
	self.diffusion = diffusion
	self.num_inference_steps = num_inference_steps

	@torch.no_grad()
	def sample(
	self,
	prompt_embeds: torch.Tensor,
	height: int = 512,
	width: int = 512,
	guidance_scale: float = 7.5,
	progress_bar: bool = True
	) -> torch.Tensor:
	"""DPM采样"""
	self.model.eval()

	# 初始化潜在表示
	latents = torch.randn(
	(1, self.model.in_channels, height // 8, width // 8),
	device=prompt_embeds.device
	)

	# 简化的DPM采样
	timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)

	if progress_bar:
	timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="DPM Sampling")
	else:
	timesteps_iter = enumerate(timesteps[:-1])

	for i, t in timesteps_iter:
	# 预测噪声
	noise_pred = self.model(latents, t.unsqueeze(0) * 999, prompt_embeds)

	# 应用引导
	if guidance_scale > 1.0:
	# 简单引导
	noise_pred = noise_pred * guidance_scale

	# DPM更新步骤
	dt = timesteps[i + 1] - t
	latents = latents + dt * noise_pred

	return latents


	class LCMSampler:
	"""LCM（潜在一致性模型）采样器，极快"""
	def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 4):
	self.model = model
	self.diffusion = diffusion
	self.num_inference_steps = num_inference_steps

	# LCM特定的参数
	self.c_skip = 1.0
	self.c_out = 1.0
	self.c_in = 1.0
	self.c_noise = 1.0

	@torch.no_grad()
	def sample(
	self,
	prompt_embeds: torch.Tensor,
	height: int = 512,
	width: int = 512,
	guidance_scale: float = 7.5,
	progress_bar: bool = True
	) -> torch.Tensor:
	"""LCM采样（极快，只需要4-8步）"""
	self.model.eval()

	# 初始化潜在表示
	latents = torch.randn(
	(1, self.model.in_channels, height // 8, width // 8),
	device=prompt_embeds.device
	)

	# LCM采样循环
	timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)

	if progress_bar:
	timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="LCM Sampling")
	else:
	timesteps_iter = enumerate(timesteps[:-1])

	for i, t in timesteps_iter:
	# LCM特定的缩放
	c_skip = self.c_skip
	c_out = self.c_out
	c_in = self.c_in
	c_noise = self.c_noise

	# 缩放输入
	scaled_latents = c_in * latents

	# 预测
	noise_pred = self.model(scaled_latents, c_noise * t.unsqueeze(0), prompt_embeds)

	# LCM更新规则
	denoised = c_skip * latents + c_out * noise_pred

	# 更新潜在表示
	dt = timesteps[i + 1] - t
	latents = denoised + dt * noise_pred

	return latents


	class SamplerFactory:
	"""采样器工厂"""
	@staticmethod
	def create_sampler(
	sampler_type: str,
	model: nn.Module,
	diffusion,
	num_inference_steps: int = 50
	):
	"""创建采样器"""
	if sampler_type == "ddim":
	return DDIMSampler(model, diffusion, num_inference_steps)
	elif sampler_type == "dpm":
	return DPMSampler(model, diffusion, num_inference_steps)
	elif sampler_type == "lcm":
	return LCMSampler(model, diffusion, num_inference_steps)
	else:
	raise ValueError(f"未知的采样器类型: {sampler_type}")


	class TextToImagePipeline:
	"""文本到图像管道"""
	def __init__(
	self,
	model: nn.Module,
	diffusion,
	text_encoder,
	vae_decoder,
	sampler_type: str = "ddim",
	device: str = "cuda"
	):
	self.model = model.to(device)
	self.diffusion = diffusion
	self.text_encoder = text_encoder
	self.vae_decoder = vae_decoder
	self.sampler_type = sampler_type
	self.device = device

	# 创建采样器
	self.sampler = SamplerFactory.create_sampler(
	sampler_type, model, diffusion
	)

	# 设置为评估模式
	self.model.eval()
	if self.vae_decoder is not None:
	self.vae_decoder.eval()

	@torch.no_grad()
	def __call__(
	self,
	prompt: str,
	negative_prompt: str = "",
	height: int = 512,
	width: int = 512,
	num_inference_steps: int = 50,
	guidance_scale: float = 7.5,
	num_images: int = 1,
	seed: Optional[int] = None,
	progress_bar: bool = True
	) -> List:
	"""生成图像"""
	# 设置随机种子
	if seed is not None:
	torch.manual_seed(seed)
	if torch.cuda.is_available():
	torch.cuda.manual_seed(seed)

	# 编码提示
	prompt_embeds = self.text_encoder.encode([prompt]).to(self.device)
	negative_prompt_embeds = None

	if negative_prompt:
	negative_prompt_embeds = self.text_encoder.encode([negative_prompt]).to(self.device)

	# 生成潜在表示
	latents = self.sampler.sample(
	prompt_embeds=prompt_embeds,
	negative_prompt_embeds=negative_prompt_embeds,
	height=height,
	width=width,
	num_images_per_prompt=num_images,
	guidance_scale=guidance_scale,
	progress_bar=progress_bar
	)

	# 解码为图像
	images = []
	for i in range(num_images):
	latent = latents[i:i+1]

	if self.vae_decoder is not None:
	image = self.vae_decoder(latent)
	else:
	# 如果没有VAE解码器，返回潜在表示
	image = latent

	images.append(image.cpu())

	return images

	def generate_grid(
	self,
	prompts: List[str],
	grid_size: Tuple[int, int] = (2, 2),
	**kwargs
	) -> torch.Tensor:
	"""生成图像网格"""
	images = []

	for prompt in prompts[:grid_size[0] * grid_size[1]]:
	image = self(prompt, **kwargs)[0]
	images.append(image)

	# 创建网格
	from torchvision.utils import make_grid
	grid = make_grid(torch.cat(images, dim=0), nrow=grid_size[1])

	return grid


	def test_sampler():
	"""测试采样器"""
	import torch.nn as nn

	# 创建模拟模型
	class MockModel(nn.Module):
	def __init__(self):
	super().__init__()
	self.in_channels = 4

	def forward(self, x, t, context):
	# 返回随机噪声
	return torch.randn_like(x)

	# 创建模拟扩散过程
	class MockDiffusion:
	def __init__(self):
	self.num_train_timesteps = 1000
	self.alphas_cumprod = torch.ones(1000)
	self.prediction_type = "epsilon"

	def extract(self, a, t, x_shape):
	return torch.ones(x_shape[0], 1, 1, 1)

	def scale_model_input(self, x, t):
	return x

	model = MockModel()
	diffusion = MockDiffusion()

	# 测试DDIM采样器
	sampler = DDIMSampler(model, diffusion, num_inference_steps=10)

	# 测试采样
	prompt_embeds = torch.randn(1, 77, 768)
	latents = sampler.sample(prompt_embeds, height=64, width=64, progress_bar=False)

	print(f"DDIM采样完成，潜在表示形状: {latents.shape}")

	return sampler, latents


	if __name__ == '__main__':
	test_sampler()