LEDITS++: Generation Steps

def ledits_core_loop_simplified(
    unet,                           # UNet模型
    scheduler,                      # 调度器
    latents,                        # 初始潜在变量 x_T
    timesteps,                      # 时间步序列
    uncond_embeddings,              # 无条件文本嵌入 (对应 "")
    edit_concepts_embeddings,       # 编辑概念的文本嵌入列表 (如 ['george clooney', 'sunglasses'])
    edit_guidance_scales,           # 每个编辑概念的引导强度 (如 [3.0, 4.0])
    edit_masks,                     # 每个编辑概念的最终掩码 (假设已经计算好)
    reverse_editing_directions,     # 每个编辑概念的方向 (如 [False, False])
    edit_weights=None,              # 每个编辑概念的权重 (可选)
    verbose=True
):
    """
    LEDITS++ 核心去噪循环的简化版本
    
    主要创新点：
    1. 并行预测多种噪声 (1次UNet调用预测1+N种噪声)
    2. 为每个编辑概念独立计算编辑向量
    3. 使用掩码进行空间定位编辑
    4. 合成多个编辑概念的最终引导
    """
    
    # 获取编辑概念数量
    num_edit_concepts = len(edit_concepts_embeddings)
    print(f"开始LEDITS++去噪循环，共有 {num_edit_concepts} 个编辑概念")
    
    # 如果没有提供权重，默认所有概念权重相等
    if edit_weights is None:
        edit_weights = [1.0] * num_edit_concepts
    
    # ==================== 核心去噪循环 ====================
    for i, timestep in enumerate(timesteps):
        if verbose:
            print(f"\n--- 时间步 {i+1}/{len(timesteps)}, t={timestep} ---")
        
        # ==================== 步骤1: 并行预测多种噪声 ====================
        print("步骤1: 并行预测多种噪声")
        
        # 准备输入：将潜在变量复制 (1 + N) 次，用于并行预测
        # 第1个用于无条件预测，后N个用于各个编辑概念
        latent_model_input = torch.cat([latents] * (1 + num_edit_concepts))
        latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)
        
        # 准备文本嵌入：无条件 + 所有编辑概念
        text_embeddings = torch.cat([uncond_embeddings] + edit_concepts_embeddings)
        
        # 一次性UNet前向传播，并行预测所有噪声
        print(f"  - 输入形状: {latent_model_input.shape}")
        print(f"  - 文本嵌入形状: {text_embeddings.shape}")
        
        # 这里是关键：一次UNet调用预测1+N种噪声
        noise_pred_all = unet(latent_model_input, timestep, encoder_hidden_states=text_embeddings).sample
        
        # 分离不同的噪声预测
        noise_predictions = noise_pred_all.chunk(1 + num_edit_concepts)
        noise_pred_uncond = noise_predictions[0]  # 无条件噪声预测
        noise_pred_edit_concepts = noise_predictions[1:]  # 各编辑概念的噪声预测
        
        print(f"  - 无条件噪声预测形状: {noise_pred_uncond.shape}")
        print(f"  - 编辑概念噪声预测数量: {len(noise_pred_edit_concepts)}")
        
        # ==================== 步骤2: 为每个编辑概念计算编辑向量和应用掩码 ====================
        print("步骤2: 为每个编辑概念计算编辑向量和应用掩码")
        
        # 存储每个概念的最终编辑向量
        masked_edit_vectors = []
        
        for c, noise_pred_edit in enumerate(noise_pred_edit_concepts):
            concept_name = f"概念{c+1}"  # 实际应用中可以传入概念名称
            print(f"  处理 {concept_name}:")
            
            # 2.1 计算原始编辑向量 (创新点2)
            edit_vector = noise_pred_edit - noise_pred_uncond
            print(f"    - 计算编辑向量: noise_pred_edit - noise_pred_uncond")
            
            # 2.2 应用编辑方向
            if reverse_editing_directions[c]:
                edit_vector = edit_vector * -1
                print(f"    - 应用反向编辑方向")
            
            # 2.3 应用引导强度缩放
            edit_vector = edit_vector * edit_guidance_scales[c]
            print(f"    - 应用引导强度: {edit_guidance_scales[c]}")
            
            # 2.4 应用掩码进行空间定位 (创新点3)
            # 这里假设edit_masks[c]已经是计算好的最终掩码
            masked_edit_vector = edit_vector * edit_masks[c]
            print(f"    - 应用空间掩码，掩码形状: {edit_masks[c].shape}")
            
            # 2.5 应用概念权重
            masked_edit_vector = masked_edit_vector * edit_weights[c]
            print(f"    - 应用概念权重: {edit_weights[c]}")
            
            masked_edit_vectors.append(masked_edit_vector)
        
        # ==================== 步骤3: 合成最终引导 ====================
        print("步骤3: 合成最终引导")
        
        # 初始化总引导为零
        total_guidance = torch.zeros_like(noise_pred_uncond)
        
        # 累加所有编辑概念的贡献
        for c, masked_vector in enumerate(masked_edit_vectors):
            total_guidance = total_guidance + masked_vector
            print(f"  - 累加概念{c+1}的编辑向量")
        
        # 计算最终的噪声预测
        final_noise_pred = noise_pred_uncond + total_guidance
        print(f"  - 最终噪声预测 = 无条件预测 + 总引导")
        
        # ==================== 步骤4: 更新潜在变量 (传统去噪步骤) ====================
        print("步骤4: 更新潜在变量")
        
        # 使用调度器进行去噪步骤: x_t -> x_{t-1}
        latents = scheduler.step(final_noise_pred, timestep, latents).prev_sample
        print(f"  - 潜在变量更新完成，新形状: {latents.shape}")
    
    print("\n==================== 去噪循环完成 ====================")
    return latents