Stable Diffusion XL 文生图

西西弗Sisyphus2025-05-12 12:30
Stable Diffusion XL 文生图

flyfish
py 复制代码
import torch
from diffusers import DiffusionPipeline, AutoencoderKL
from PIL import Image

# 定义模型路径为常量
BASE_MODEL_PATH = "/media/stable-diffusion-xl-base-1___0/"
REFINER_MODEL_PATH = "/media/stable-diffusion-xl-refiner-1___0/"
VAE_MODEL_PATH = "/media/sdxl-vae-fp16-fix/"

# 加载 VAE 模型
vae = AutoencoderKL.from_pretrained(VAE_MODEL_PATH, torch_dtype=torch.float16)

# 加载基础模型和细化器模型
base_pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL_PATH, vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
refiner_pipe = DiffusionPipeline.from_pretrained(
    REFINER_MODEL_PATH, vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
)

# 将模型移动到 GPU
base_pipe.to("cuda")
refiner_pipe.to("cuda")

# 设置推理步骤和其他参数
n_steps = 50
high_noise_frac = 0.7

# 自定义生成图像的分辨率
width = 1920  # 宽度
height = 1080  # 高度

# 输入提示词
prompt = "A beautiful and dreamy world with floating islands, glowing rivers, and magical creatures"
negative_prompt = "ugly, blurry, distorted, disfigured, low quality"

# 生成图像（使用自定义分辨率）
latent_image = base_pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
    width=width,
    height=height
).images


# 细化图像并保存最终结果
final_image = refiner_pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=latent_image
).images[0]

final_image.save("refined_generated_image.png")
print("细化模型图像已保存为 refined_generated_image.png")