文本到图像生成：DALL-E, Imagen

# 文本到图像生成的基本抽象
import torch
from typing import Optional

class TextToImageModel:
    """文本到图像生成模型的统一接口"""

    def __init__(self, text_encoder, generator, image_decoder):
        self.text_encoder = text_encoder    # 将文本编码为向量
        self.generator = generator          # 核心生成网络
        self.image_decoder = image_decoder  # 解码为像素图像

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        guidance_scale: float = 7.5,
        num_inference_steps: int = 50,
        seed: Optional[int] = None,
    ) -> torch.Tensor:
        # Step 1: 编码文本
        text_embed = self.text_encoder(prompt)

        # Step 2: 初始化随机噪声
        if seed is not None:
            torch.manual_seed(seed)
        latent = torch.randn(1, 4, 64, 64, device=text_embed.device)

        # Step 3: 迭代生成（扩散过程或自回归过程）
        for step in range(num_inference_steps):
            noise_pred = self.generator(latent, step, text_embed)
            # 分类器自由引导（CFG）
            noise_pred_uncond = self.generator(latent, step, torch.zeros_like(text_embed))
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
            latent = self._step(latent, noise_pred, step)

        # Step 4: 解码为图像
        image = self.image_decoder(latent)
        return image

    def _step(self, latent, noise_pred, step):
        """DDIM 或 Euler 采样单步更新"""
        alpha = 1.0 - step / 1000.0
        return latent - 0.1 * noise_pred * alpha

# 评估生成图像与文本的一致性
import torch
import clip

def clip_score(image: torch.Tensor, prompt: str, model, preprocess) -> float:
    """计算 CLIP Score：图像-文本相似度"""
    # 预处理图像
    image_input = preprocess(image).unsqueeze(0)

    # 编码文本
    text_tokens = clip.tokenize([prompt])

    with torch.no_grad():
        image_features = model.encode_image(image_input)
        text_features = model.encode_text(text_tokens)

    # 归一化并计算余弦相似度
    image_features = image_features / image_features.norm(dim=-1, keepdim=True)
    text_features = text_features / text_features.norm(dim=-1, keepdim=True)
    similarity = (image_features @ text_features.T).item()

    return similarity

# 测试
model, preprocess = clip.load("ViT-L/14", device="cuda")
image = torch.rand(3, 224, 224)
score = clip_score(image, "一只在草地上奔跑的柯基犬", model, preprocess)
print(f"CLIP Score: {score:.4f}")  # 越接近 1 表示越匹配

# 加载 CLIP 模型并计算图文相似度
import torch
import clip
from PIL import Image

# 加载预训练模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-L/14", device=device)

# 准备输入
image = preprocess(Image.open("dog.jpg")).unsqueeze(0).to(device)
texts = clip.tokenize([
    "一只棕色的狗在草地上玩耍",
    "一只猫坐在窗台上",
    "一辆红色的跑车",
]).to(device)

# 编码
with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(texts)

    # 归一化
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)

    # 计算相似度
    similarity = (image_features @ text_features.T).softmax(dim=-1)
    print(f"狗的相似度: {similarity[0, 0].item():.4f}")
    print(f"猫的相似度: {similarity[0, 1].item():.4f}")
    print(f"车的相似度: {similarity[0, 2].item():.4f}")

# CLIP 对比损失函数的简化实现
import torch
import torch.nn as nn
import torch.nn.functional as F

class CLIPLoss(nn.Module):
    """CLIP 的 InfoNCE 对比损失"""

    def __init__(self, temperature: float = 0.07):
        super().__init__()
        self.logit_scale = nn.Parameter(torch.tensor(1.0 / temperature))

    def forward(
        self,
        image_features: torch.Tensor,  # [batch_size, dim]
        text_features: torch.Tensor,   # [batch_size, dim]
    ) -> torch.Tensor:
        # 归一化特征
        image_features = F.normalize(image_features, dim=-1)
        text_features = F.normalize(text_features, dim=-1)

        # 计算相似度矩阵
        logits_per_image = self.logit_scale.exp() * image_features @ text_features.T
        logits_per_text = logits_per_image.T

        # 标签：对角线为正样本
        batch_size = image_features.size(0)
        labels = torch.arange(batch_size, device=image_features.device)

        # 双向交叉熵损失
        loss_i = F.cross_entropy(logits_per_image, labels)
        loss_t = F.cross_entropy(logits_per_text, labels)
        return (loss_i + loss_t) / 2

# 模拟训练
loss_fn = CLIPLoss()
img_feat = torch.randn(32, 512)
txt_feat = torch.randn(32, 512)
loss = loss_fn(img_feat, txt_feat)
print(f"CLIP Loss: {loss.item():.4f}")

# 简化版 DALL-E 架构：dVAE + Transformer
import torch
import torch.nn as nn
import torch.nn.functional as F

class VectorQuantizer(nn.Module):
    """VQ-VAE 的向量量化层"""

    def __init__(self, num_embeddings=8192, embedding_dim=256):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.embedding.weight.data.uniform_(-1/num_embeddings, 1/num_embeddings)

    def forward(self, z: torch.Tensor):
        # z: [B, C, H, W] -> [B, H*W, C]
        z_flat = z.flatten(2).permute(0, 2, 1)
        # 计算与所有 codebook 向量的距离
        distances = torch.cdist(z_flat, self.embedding.weight)
        indices = torch.argmin(distances, dim=-1)  # [B, H*W]
        quantized = self.embedding(indices)        # [B, H*W, C]
        return quantized, indices

class DALLETransformer(nn.Module):
    """DALL-E 的自回归 Transformer"""

    def __init__(self, vocab_size=8192, text_vocab=50257, d_model=768, n_layers=24):
        super().__init__()
        self.text_embed = nn.Embedding(text_vocab, d_model)
        self.image_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Parameter(torch.randn(1, 2048, d_model) * 0.02)
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead=12, dim_feedforward=3072)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, text_tokens, image_tokens):
        text_emb = self.text_embed(text_tokens)
        image_emb = self.image_embed(image_tokens)
        combined = torch.cat([text_emb, image_emb], dim=1)
        combined = combined + self.pos_embed[:, :combined.size(1)]
        hidden = self.transformer(combined.transpose(0, 1))
        # 只预测图像部分
        img_hidden = hidden[-image_tokens.size(1):]
        return self.lm_head(img_hidden.transpose(0, 1))

# DALL-E 风格的条件图像生成（自回归采样）
import torch

def dalle_generate(model, text_tokens, max_len=1024, top_k=256, temperature=1.0):
    """自回归生成图像 token 序列"""
    model.eval()
    image_tokens = []

    with torch.no_grad():
        for i in range(max_len):
            # 逐步生成
            current_image = torch.tensor([image_tokens]) if image_tokens else torch.empty((1, 0), dtype=torch.long)
            logits = model(text_tokens, current_image.to(text_tokens.device))

            # 取最后一个 token 的 logits
            next_logits = logits[:, -1, :] / temperature

            # top-k 采样
            if top_k > 0:
                topk_values, topk_indices = torch.topk(next_logits, top_k)
                probs = F.softmax(topk_values, dim=-1)
                sample_idx = torch.multinomial(probs, 1)
                next_token = topk_indices.gather(1, sample_idx)
            else:
                probs = F.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, 1)

            image_tokens.append(next_token.item())

    return torch.tensor(image_tokens).view(32, 32)  # 重排为 32x32 网格

# 使用示例
text_tokens = torch.randint(0, 50257, (1, 64))  # 模拟 BPE 编码
# model = DALLETransformer()
# tokens = dalle_generate(model, text_tokens)
print("生成了 1024 个图像 token，形状: 32x32")

# Imagen 级联扩散管道的简化实现
import torch
import torch.nn as nn

class ImagenPipeline:
    """Imagen 级联扩散模型管道"""

    def __init__(self, text_encoder, base_diffusion, sr_model_1, sr_model_2):
        self.text_encoder = text_encoder   # T5-XXL
        self.base_diffusion = base_diffusion  # 64x64
        self.sr_model_1 = sr_model_1       # 64->256
        self.sr_model_2 = sr_model_2       # 256->1024

    @torch.no_grad()
    def generate(self, prompt: str, guidance_scale: float = 7.5, seed: int = 42):
        # Step 1: 文本编码（冻结）
        text_embed = self.text_encoder.encode(prompt)

        # Step 2: 64x64 基础生成
        torch.manual_seed(seed)
        low_res = self.base_diffusion.sample(
            cond=text_embed,
            guidance_scale=guidance_scale,
            steps=100,
        )  # [1, 3, 64, 64]

        # Step 3: 第一次超分辨率 64 -> 256
        mid_res = self.sr_model_1.sample(
            low_res_image=low_res,
            cond=text_embed,
            guidance_scale=guidance_scale,
            steps=50,
        )  # [1, 3, 256, 256]

        # Step 4: 第二次超分辨率 256 -> 1024
        high_res = self.sr_model_2.sample(
            low_res_image=mid_res,
            cond=text_embed,
            guidance_scale=guidance_scale,
            steps=50,
        )  # [1, 3, 1024, 1024]

        return high_res

# Dynamic Thresholding - Imagen 的关键技术
import torch

def dynamic_thresholding(
    samples: torch.Tensor,
    percentile: float = 0.995,
    clamp_min: float = -1.0,
    clamp_max: float = 1.0,
) -> torch.Tensor:
    """Imagen 的动态阈值化技术

    解决 CFG 引导过强导致的色彩饱和和伪影问题。
    核心思想：根据当前样本的绝对值分布动态调整裁剪范围，
    而不是使用固定的 [-1, 1] 范围。
    """
    device = samples.device
    dtype = samples.dtype

    # 展平并计算绝对值的分位数
    s = samples.flatten(1).abs()
    s_quantile, _ = torch.quantile(s, percentile, dim=1)  # [batch]
    s_quantile = torch.clamp(s_quantile, min=1.0)  # 确保最小范围为 1.0

    # 归一化到 [-1, 1]
    s_quantile = s_quantile.view(-1, 1, 1, 1)
    normalized = samples / s_quantile

    # 最终裁剪
    return torch.clamp(normalized, clamp_min, clamp_max).to(dtype)

# 测试
noisy_pred = torch.randn(1, 3, 64, 64) * 3.0  # 模拟高 CFG 导致的溢出
cleaned = dynamic_thresholding(noisy_pred)
print(f"原始范围: [{noisy_pred.min():.2f}, {noisy_pred.max():.2f}]")
print(f"阈值化后: [{cleaned.min():.2f}, {cleaned.max():.2f}]")

# DALL-E 2 风格的先验网络
import torch
import torch.nn as nn

class DiffusionPrior(nn.Module):
    """DALL-E 2 的先验网络：将 CLIP 文本嵌入映射到 CLIP 图像嵌入"""

    def __init__(self, text_dim: int = 768, image_dim: int = 768, hidden_dim: int = 1024):
        super().__init__()
        # 简单的 MLP 先验（实际中使用扩散过程）
        self.text_proj = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, image_dim),
        )

    def forward(self, text_embed: torch.Tensor, noise: torch.Tensor = None) -> torch.Tensor:
        """将文本嵌入投影到图像嵌入空间"""
        # 简化版：直接投影（实际使用扩散逆过程）
        projected = self.text_proj(text_embed)
        if noise is not None:
            projected = projected + noise
        return projected

# 使用示例
prior = DiffusionPrior(text_dim=768, image_dim=768)
text_embed = torch.randn(1, 768)  # 来自 CLIP 文本编码器
image_embed = prior(text_embed)
print(f"CLIP 图像嵌入形状: {image_embed.shape}")  # [1, 768]

# DALL-E 3 风格的 prompt 扩展
import openai

class PromptExpander:
    """模拟 DALL-E 3 的 GPT-4 prompt 扩展机制"""

    def __init__(self, api_key: str = None):
        self.client = openai.OpenAI(api_key=api_key)
        self.system_prompt = """你是一个专业的视觉描述专家。
用户会给你一个简短的图像生成 prompt，你需要将其扩展为
一个详细的描述，包含：主体、背景、构图、风格、色彩、光影。
扩展后的描述应该具体、生动，适合文本到图像模型使用。"""

    def expand(self, short_prompt: str, style: str = None) -> str:
        """将简短 prompt 扩展为详细描述"""
        style_hint = f"\n风格偏好: {style}" if style else ""
        response = self.client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": f"扩展以下 prompt：{short_prompt}{style_hint}"}
            ],
            temperature=0.7,
            max_tokens=300,
        )
        return response.choices[0].message.content

# 使用示例
expander = PromptExpander()
expanded = expander.expand(
    short_prompt="一只猫在太空",
    style="赛博朋克，霓虹色调"
)
print(f"扩展后的 prompt:\n{expanded}")

# 计算 FID 分数
import numpy as np
from scipy import linalg
import torch
import torchvision.models as models
from torchvision import transforms

class FIDEvaluator:
    """Frechet Inception Distance 计算"""

    def __init__(self, device="cuda"):
        self.device = device
        # 加载 Inception-V3 并移除分类头
        self.model = models.inception_v3(
            weights=models.Inception_V3_Weights.IMAGENET1K_V1
        )
        self.model.fc = torch.nn.Identity()  # 移除分类头
        self.model.eval().to(device)

        # 预处理
        self.transform = transforms.Compose([
            transforms.Resize((299, 299)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225]),
        ])

    def extract_features(self, images: torch.Tensor) -> np.ndarray:
        """提取 Inception-V3 特征"""
        features = []
        with torch.no_grad():
            for img in images:
                feat = self.model(img.unsqueeze(0).to(self.device))
                features.append(feat.cpu().numpy())
        return np.concatenate(features, axis=0)

    @staticmethod
    def calculate_fid(
        mu1: np.ndarray, sigma1: np.ndarray,
        mu2: np.ndarray, sigma2: np.ndarray,
    ) -> float:
        """计算两个高斯分布之间的 Frechet 距离"""
        diff = mu1 - mu2
        # 协方差矩阵的平方根（使用 matrix square root）
        covmean, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False)
        if np.iscomplexobj(covmean):
            covmean = covmean.real
        fid = diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(covmean)
        return float(fid)

    def evaluate(self, real_images: torch.Tensor, gen_images: torch.Tensor) -> float:
        real_feats = self.extract_features(real_images)
        gen_feats = self.extract_features(gen_images)
        mu1, sigma1 = real_feats.mean(0), np.cov(real_feats, rowvar=False)
        mu2, sigma2 = gen_feats.mean(0), np.cov(gen_feats, rowvar=False)
        return self.calculate_fid(mu1, sigma1, mu2, sigma2)

# 综合评估：FID + CLIP Score + PickScore
import torch

class ComprehensiveEvaluator:
    """综合评估文本到图像生成质量"""

    def __init__(self, clip_model, inception_model, device="cuda"):
        self.clip_model = clip_model
        self.inception_model = inception_model
        self.device = device

    def batch_evaluate(
        self,
        prompts: list[str],
        generated_images: torch.Tensor,
        real_images: torch.Tensor = None,
    ) -> dict:
        results = {}

        # 1. CLIP Score（图文一致性）
        clip_scores = []
        with torch.no_grad():
            text_embeds = self.clip_model.encode_text(
                torch.stack([torch.tensor(clip.tokenize(p)) for p in prompts])
            )
            image_embeds = self.clip_model.encode_image(generated_images)
            text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
            image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
            clip_scores = (image_embeds @ text_embeds.T).diag().tolist()

        results["mean_clip_score"] = sum(clip_scores) / len(clip_scores)
        results["clip_scores"] = clip_scores

        # 2. FID（图像质量，需要真实图像参考）
        if real_images is not None:
            pass  # fid = fid_eval.evaluate(real_images, generated_images)
            # results["fid"] = fid

        # 3. 多样性（pairwise CLIP 距离）
        image_embeds_norm = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
        pairwise_sim = image_embeds_norm @ image_embeds_norm.T
        # 去除对角线（自身相似度=1）
        mask = ~torch.eye(pairwise_sim.size(0), dtype=torch.bool)
        diversity = 1.0 - pairwise_sim[mask].mean().item()
        results["diversity"] = diversity

        return results

# 使用示例
evaluator = ComprehensiveEvaluator(clip_model=None, inception_model=None)
prompts = ["一只狗", "一只猫", "一朵花"]
gen_images = torch.randn(3, 3, 224, 224)
# results = evaluator.batch_evaluate(prompts, gen_images)

# diffusers 实战：自定义文本到图像 Pipeline
from diffusers import (
    StableDiffusionPipeline,
    EulerAncestralDiscreteScheduler,
    DPMSolverMultistepScheduler,
)
import torch
from PIL import Image

class CustomImageGenerator:
    """基于 diffusers 的自定义图像生成器"""

    def __init__(self, model_id: str = "runwayml/stable-diffusion-v1-5"):
        # 加载基础 Pipeline
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            safety_checker=None,  # 生产环境建议保留
            use_safetensors=True,
        )

        # 使用更快的采样器
        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            self.pipe.scheduler.config
        )
        self.pipe.to("cuda")

    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 512,
        height: int = 512,
        guidance_scale: float = 7.5,
        num_inference_steps: int = 30,
        seed: int = None,
        num_images: int = 1,
    ) -> list[Image.Image]:
        """生成图像"""
        generator = None
        if seed is not None:
            generator = torch.Generator(device="cuda").manual_seed(seed)

        images = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
            num_images_per_prompt=num_images,
        ).images

        return images

    def generate_batch(
        self,
        prompts: list[str],
        **kwargs,
    ) -> list[list[Image.Image]]:
        """批量生成多张图像"""
        all_images = []
        for prompt in prompts:
            images = self.generate(prompt=prompt, **kwargs)
            all_images.append(images)
        return all_images

# 使用示例
gen = CustomImageGenerator()
images = gen.generate(
    prompt="一只赛博朋克风格的猫，霓虹灯光，未来城市背景",
    negative_prompt="模糊, 低质量, 变形的手",
    guidance_scale=8.0,
    num_inference_steps=40,
    seed=42,
)
images[0].save("cyberpunk_cat.png")

# diffusers 进阶：LoRA 微调与自定义采样
from diffusers import StableDiffusionPipeline, DDIMScheduler
from peft import LoraConfig, inject_adapter_in_model
import torch

class LoRAImageGenerator:
    """使用 LoRA 进行高效微调的生成器"""

    def __init__(self, base_model: str, lora_path: str = None):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            base_model,
            torch_dtype=torch.float16,
            use_safetensors=True,
        )
        self.pipe.scheduler = DDIMScheduler.from_config(
            self.pipe.scheduler.config,
            timestep_spacing="trailing",
            prediction_type="epsilon",
        )
        self.pipe.to("cuda")

        # 加载 LoRA 权重
        if lora_path:
            self.pipe.load_lora_weights(lora_path)
            print(f"已加载 LoRA 权重: {lora_path}")

    def generate_with_cfg(
        self,
        prompt: str,
        negative_prompt: str = "",
        guidance_scale: float = 7.5,
        eta: float = 0.0,  # DDIM eta 参数
    ) -> torch.Tensor:
        """带分类器自由引导的 DDIM 采样"""
        # 编码 prompt
        text_input = self.pipe.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.pipe.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        ).to(self.pipe.device)

        text_embeds = self.pipe.text_encoder(text_input.input_ids)[0]

        # 编码 negative prompt
        if negative_prompt:
            uncond_input = self.pipe.tokenizer(
                negative_prompt,
                padding="max_length",
                max_length=self.pipe.tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            ).to(self.pipe.device)
            uncond_embeds = self.pipe.text_encoder(uncond_input.input_ids)[0]
            text_embeds = torch.cat([uncond_embeds, text_embeds])

        # 初始化噪声
        latent = torch.randn(1, 4, 64, 64, device=self.pipe.device, dtype=torch.float16)

        # DDIM 采样循环
        self.pipe.scheduler.set_timesteps(50)
        for t in self.pipe.scheduler.timesteps:
            latent_input = torch.cat([latent] * 2) if negative_prompt else latent
            noise_pred = self.pipe.unet(latent_input, t, text_embeds).sample

            if negative_prompt:
                noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (
                    noise_pred_cond - noise_pred_uncond
                )

            latent = self.pipe.scheduler.step(noise_pred, t, latent, eta=eta).prev_sample

        # 解码
        latent = 1 / self.pipe.vae.config.scaling_factor * latent
        image = self.pipe.vae.decode(latent).sample
        image = (image / 2 + 0.5).clamp(0, 1)
        return image