如何在本地运行Juggernaut模型

Question

我想在我的电脑上用Python运行经过微调的稳定扩散模型。比如说这个叫“Juggernaut”的模型：https://huggingface.co/RunDiffusion/Juggernaut-XL-v9

这是我的代码（它在使用stable-diffusion-xl-base-1.0时可以正常工作）：

import random
from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
import torch
import gc
import time

# for cleaning memory
gc.collect()
torch.cuda.empty_cache()

start_time = time.time()

model = "RunDiffusion/Juggernaut-XL-v9"
pipe = DiffusionPipeline.from_pretrained(
    model,
    torch_dtype=torch.float16,
)

pipe.to("cuda")

prompt = ("a portrait of male as a knight in middle ages, masculine looking, battle in the background, sharp focus, highly detailed, movie-style lighting, shadows")
seed = random.randint(0, 2**32 - 1)

generator = torch.Generator("cuda").manual_seed(seed)
image = pipe(prompt=prompt, generator=generator, num_inference_steps=1)
image = image.images[0]
image.save(f"output_images/{seed}.png")

end_time = time.time()

total_time = end_time - start_time
minutes = int(total_time // 60) 
seconds = int(total_time % 60) 

print(f"Took: {minutes} min {seconds} sec")
print(f"Saved to output_images/{seed}.png")

但是我遇到了这个错误：

OSError: 错误，没有找到名为pytorch_model.bin、tf_model.h5、model.ckpt.index或flax_model.msgpack的文件在目录中

可能是因为Python和CUDA的版本问题。我正在降低我的库的版本：

Python 3.9.0

PyTorch: 2.2.0+cu118

CUDA : 11.8

Diffusers: 0.26.3

Transformers: 4.38.1

Answer 1

这段代码的意思是：我们要创建一个叫做“pipe”的东西，它是通过一个叫“DiffusionPipeline”的工具来生成的。这个工具会从一个预先准备好的模型中获取信息。我们还指定了一些参数，比如使用16位浮点数（torch.float16）来提高计算效率，选择一种叫“fp16”的变体，并且启用一种叫“safetensors”的安全选项。