pix2pix

import PIL
from PIL import Image, ImageOps
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
import os

model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
    safety_checker=None
)
pipe.to("cuda")
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

local_image_path = "input.jpg"

def load_local_image(image_path):
    try:
        image = Image.open(image_path)
        image = ImageOps.exif_transpose(image)
        image = image.convert("RGB")
        return image
    except Exception as e:
        print(f"无法加载图片: {e}")
        return None

image = load_local_image(local_image_path)

if image is None:
    raise ValueError("图片加载失败，请检查路径是否正确。")

prompt = "turn him into cyborg"
images = pipe(
    prompt, 
    image=image, 
    num_inference_steps=10, 
    image_guidance_scale=1
).images

current_dir = os.getcwd()
print(f"当前工作目录: {current_dir}")

output_image_path = os.path.join(current_dir, "output_cyborg.png")
images[0].save(output_image_path)
print(f"生成的图片已保存到: {output_image_path}")

images[0].show()