import torch
from diffusers import StableDiffusion3Pipeline

# --- 1. 初始化模型和调度器（采样器） ---
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", 
    torch_dtype=torch.float16
)

# 获取核心 MMDiT 模型组件
# === 关键修正：使用 pipe.components["unet"] 访问去噪模型 ===
MMDIT_MODEL = pipe.components["unet"] 

pipe.to("cuda")

# 文本提示和参数
prompt = "A majestic castle on a floating island, photorealistic, 4k"
negative_prompt = "" 

# 引导尺度（CFG Scale）
guidance_scale = 7.0 
num_inference_steps = 20
pipe.scheduler.set_timesteps(num_inference_steps, device=pipe.device) # 设置步数

# --- 2. 准备嵌入 (SD3多编码器修正) ---
prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
    prompt=prompt,
    prompt_2=prompt, 
    prompt_3=prompt, 
    negative_prompt=negative_prompt,
    negative_prompt_2=negative_prompt, 
    negative_prompt_3=negative_prompt,
    device=pipe.device,
    do_classifier_free_guidance=True, 
)

# 堆叠嵌入
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)

# --- 3. 采样循环定制 ---

# 初始化噪声 (Latent)
LATENT_CHANNELS = 16 
latents = torch.randn(
    (1, LATENT_CHANNELS, 1024 // 8, 1024 // 8), 
    generator=None,
    device=pipe.device,
    dtype=pipe.dtype,
)

print("开始定制采样循环...")
for i, t in enumerate(pipe.scheduler.timesteps):
    # 扩展 latents 以同时计算条件和无条件预测
    latent_model_input = torch.cat([latents] * 2) 
    
    # UNet (实际是 MMDiT) 预测：v_uncond, v_cond
    # === 关键修正：使用 MMDIT_MODEL 变量 (即 pipe.components["unet"]) ===
    model_output = MMDIT_MODEL( 
        latent_model_input, 
        t, 
        encoder_hidden_states=prompt_embeds, 
        pooled_projections=pooled_prompt_embeds, 
        return_dict=False
    )[0]

    # 分割预测结果
    v_uncond, v_cond = model_output.chunk(2)
    
    # --- 您的定制逻辑 ---
    v_guided = v_cond 

    # ODE 求解器更新
    latents = pipe.scheduler.step(v_guided, t, latents, return_dict=False)[0]

# --- 4. 解码并保存结果 ---
print("解码中...")
latents = 1 / pipe.vae.config.scaling_factor * latents 
image = pipe.vae.decode(latents, return_dict=False)[0]
image = pipe.image_processor.postprocess(image.detach().cpu().float(), output_type="pil")
image.save("sd3_conditional_only_ode_result.png")

print("采样完成，结果已保存至 sd3_conditional_only_ode_result.png")