import PIL
import requests
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline

# 1. 指定你的checkpoint-1000所在的本地路径
# 例如：如果checkpoint保存在 "/root/training_results/checkpoint-1000"
checkpoint_path = "/root/autodl-tmp/my_self/checkpoint-100"  # <- 替换为实际路径

# 2. 从本地checkpoint加载管道
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    checkpoint_path,  # 本地checkpoint目录
    torch_dtype=torch.float16,
    safety_checker=None  # 可选：禁用安全检查器（如果训练时未启用）
).to("cuda")

# 3. 生成器设置（保持不变）
generator = torch.Generator("cuda").manual_seed(0)

# 4. 加载本地图像（修改部分）
def load_local_image(image_path):
    """从本地路径加载图像并预处理"""
    image = PIL.Image.open(image_path)  # 直接打开本地文件
    image = PIL.ImageOps.exif_transpose(image)  # 处理图像旋转信息
    image = image.convert("RGB")  # 转换为RGB格式
    return image

# 替换为你的本地图像路径（例如："/root/test_images/person.jpg"）
local_image_path = "/root/autodl-tmp/my_model/inference.png"  # <- 替换为你的本地图像路径
image = load_local_image(local_image_path)

# 5. 推理参数（根据你的训练数据调整）
prompt = "smile"  # 你的训练任务是"smile"，保持一致
num_inference_steps = 20
image_guidance_scale = 1.5
guidance_scale = 10

# 6. 生成编辑后的图像
edited_image = pipe(
    prompt,
    image=image,
    num_inference_steps=num_inference_steps,
    image_guidance_scale=image_guidance_scale,
    guidance_scale=guidance_scale,
    generator=generator,
).images[0]

def save_generated_image(image, save_dir, save_filename="generated_smile.png"):
    # 自动创建保存文件夹（避免手动建文件夹的麻烦）
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)  # exist_ok=True 防止文件夹已存在时报错
    
    # 拼接完整保存路径
    save_path = os.path.join(save_dir, save_filename)
    
    # 保存图像（PNG格式支持透明，若要存JPG，将文件名后缀改为 .jpg 即可）
    image.save(save_path)
    
    # 打印路径，方便直接找到文件
    print(f"\n生成图像已保存 → {save_path}")

# 配置保存参数（替换为你的本地路径）
save_directory = "/root/autodl-tmp/my_model"  # 例：/root/results/smile_outputs
save_filename = "ckpt-100.png"  # 自定义文件名，避免重复覆盖

# 执行保存
save_generated_image(edited_image, save_directory, save_filename)