为什么需要本地模型

最近大模型比较火，豆包、极梦、元宝等都可以根据提示词，为我们生成各种各样的图片。

但是，对于提示词，以及生成的内容，这些平台却无一例外地都加了很多很多层的限制。

首先，对于我们的提示词，有一层很严格的过滤。我们一不小心，就会返回一个“您的输入不符合平台规则，生成失败”。

其次，对于生成出来的内容，仍然有不知道多少过滤。我们经常看见，生成出来了，唰一下，没看清，又返回生成失败了。

如果把大模型部署在本地，我们就可以随心所欲地生成了。

本地大模型生成照片有多简单

不得不说，开源是真好。

使用开源的代码以及模型，在本地部署并且运行，只需要几行代码。

比如，我们要使用"runwayml/stable-diffusion-v1-5"这个模型，生成照片，只需要三行代码：

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
prompt = "lovely girl, sexy, hot"
image = pipeline(prompt).images[0]

如我们通过matplotlib来显示：

import matplotlib.pyplot as plt

plt.figure('image')
plt.imshow(image)

因为代码过于短小，只要是Python程序员，上面的代码可能都是秒懂。

即，我们需要安装diffusers库，就可以使用diffusers的各种pipeline的from_pretrained方法来加载模型。

from_pretrained方法会自动下载模型，可以根据huggingface 列表选择任何模型，或者通过diffusers 的文档，定制代码。

去除限制，随心所欲

按照默认设置，如果我们生成的图片，触发了NSFW，diffusers会对图片置黑，同时在终端返回一段：

Potential NSFW content was detected in one or more images. A black image will be returned instead.
Try again with a different prompt and/or seed.

网上有人说需要把stable_diffusion目录的safety_checker.py里面这一块直接注释掉：

  
87         for idx, has_nsfw_concept in enumerate(has_nsfw_concepts):  
88             if has_nsfw_concept:  
89                 if torch.is_tensor(images) or torch.is_tensor(images[0]):  
90                     images[idx] = torch.zeros_like(images[idx])  # black image  
91                 else:  
92                     images[idx] = np.zeros(images[idx].shape)  # black image  
93    
94         if any(has_nsfw_concepts):  
95             logger.warning(  
96                 "Potential NSFW content was detected in one or more images. A black image will be returned instead."  
97                 " Try again with a different prompt and/or seed."  
98             )

其实是不对的。

对于这么成熟的库，一层检验不至于连个开关都没有。

对于这种检查，关掉的方法非常简单，只要把相应的回调函数置空：

pipeline.safety_checker = lambda images, clip_input: (images, None)

或者加载的时候：

pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", requires_safety_checker=False)

汉语支持

这么使用模型，提示词得是英语。

我们可以再通过两行代码，加入一个本地的翻译模型：

from transformers import pipeline as trans_pipeline

translator = trans_pipeline("translation_zh_to_en", model="Helsinki-NLP/opus-mt-zh-en")

def zh_to_en(text):
    return translator(text)[0]['translation_text']

合并到一起，包装一个统一点儿的函数：

from diffusers import DiffusionPipeline
from transformers import pipeline as trans_pipeline

pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", requires_safety_checker=False)
pipeline.safety_checker = lambda images, clip_input: (images, None)
translator = trans_pipeline("translation_zh_to_en", model="Helsinki-NLP/opus-mt-zh-en")

# 两个参数：
# 1. 汉语提示词
# 2. 生成的照片数量，默认是1。
def gen_images(prompt, num=1):
    images = []
    text = translator(prompt)[0]['translation_text']
    print("translate to: [{}]".format(text))
    
    for i in range(num):
        image = pipeline(text).images[0]
        image.save("/tmp/image-{}.png".format(i))
        images.append(image)

    return images

本地下载

如果网络不好，需要手动下载模型，可以把模型文件放在$HOME/.cache/huggingface/hub目录里。

规则是前面加models–，然后再把/替换成–。

如：runwayml/stable-diffusion-v1-5，下载的目录就应该是$HOME/.cache/huggingface/hub/models–runwayml–stable-diffusion-v1-5 。