CogView4 文本生成图像

flyfish

基于 CogView4Pipeline 的图像生成程序，其主要目的是依据 JSON 文件里的文本提示信息来生成图像，并且把生成的图像保存到指定文件夹。
JSON 文件格式

[{"prompt": "your first prompt"},{"prompt": "your second prompt"}
]

从源代码安装diffusers库

pip install git+https://github.com/huggingface/diffusers.git

分辨率: 长宽均需满足 512px - 2048px 之间，需被32整除, 并保证最大像素数不超过 2^21 px。
精度: BF16 / FP32 (不支持FP16，会出现溢出导致纯黑图片)

import json
import torch
import os
from datetime import datetime
import random
import string
import argparse
from diffusers import CogView4Pipeline# 单例模式配置类
class ConfigSingleton:_instance = Nonedef __new__(cls):if cls._instance is None:cls._instance = super().__new__(cls)parser = argparse.ArgumentParser(description='Image generation configuration')parser.add_argument('--model_path', default='/media/models/ZhipuAI/CogView4-6B/',help='Path to the model')parser.add_argument('--height', type=int, default=1024, help='Height of the generated images')parser.add_argument('--width', type=int, default=1024, help='Width of the generated images')parser.add_argument('--guidance_scale', type=float, default=3.5,help='Guidance scale for image generation')parser.add_argument('--num_inference_steps', type=int, default=50,help='Number of inference steps')parser.add_argument('--num_images_per_prompt', type=int, default=10,help='Number of images to generate per prompt')parser.add_argument('--json_file', default='DonQuixotedelaMancha.json',help='Path to the JSON file containing prompts')parser.add_argument('--save_folder', default='generated_images',help='Folder to save the generated images')args = parser.parse_args()cls._instance.config = {"save_folder": args.save_folder,"filename_timestamp_format": "%Y%m%d_%H%M%S","random_char_count": 6,"model_path": args.model_path,"height": args.height,"width": args.width,"guidance_scale": args.guidance_scale,"num_inference_steps": args.num_inference_steps,"num_images_per_prompt": args.num_images_per_prompt,"json_file": args.json_file}os.makedirs(cls._instance.config["save_folder"], exist_ok=True)return cls._instancedef get_config(self):return self.config# 数据加载器类
class DataLoader:def __init__(self, config):self.config = configdef load_data(self):try:with open(self.config["json_file"], 'r') as f:return json.load(f)except FileNotFoundError:print(f"Error: The JSON file '{self.config['json_file']}' was not found.")return []except json.JSONDecodeError:print(f"Error: Failed to decode the JSON file '{self.config['json_file']}'.")return []# 模型管道初始化类
class PipelineInitializer:def __init__(self, config):self.config = configdef initialize_pipeline(self):pipe = CogView4Pipeline.from_pretrained(self.config["model_path"],torch_dtype=torch.bfloat16)# 将模型移到 GPU 上pipe = pipe.to("cuda")pipe.vae.enable_slicing()pipe.vae.enable_tiling()return pipe# 文件名生成器类
class FilenameGenerator:def __init__(self, config):self.config = configdef generate_unique_filename(self):timestamp = datetime.now().strftime(self.config["filename_timestamp_format"])random_chars = ''.join(random.choices(string.ascii_letters + string.digits, k=self.config["random_char_count"]))return f"{timestamp}_{random_chars}.png"# 图像生成和保存类
class ImageGeneratorAndSaver:def __init__(self, pipe, config):self.pipe = pipeself.config = configself.filename_generator = FilenameGenerator(config)def generate_and_save_images(self, data):# 使用 CUDA 生成器generator = torch.Generator("cuda")for i, item in enumerate(data):prompt = item["prompt"]try:result = self.pipe(prompt=prompt,guidance_scale=self.config["guidance_scale"],num_images_per_prompt=self.config["num_images_per_prompt"],num_inference_steps=self.config["num_inference_steps"],width=self.config["width"],height=self.config["height"],generator=generator)for j in range(len(result.images)):filename = self.filename_generator.generate_unique_filename()file_path = os.path.join(self.config["save_folder"], filename)result.images[j].save(file_path)print(f"Saved {file_path} for prompt: {prompt}")except Exception as e:print(f"Error generating image for prompt '{prompt}': {e}")# 主类，协调各个组件
class ImageGenerationManager:def __init__(self):self.config = ConfigSingleton().get_config()self.data_loader = DataLoader(self.config)self.pipeline_initializer = PipelineInitializer(self.config)def run(self):data = self.data_loader.load_data()if not data:returnpipe = self.pipeline_initializer.initialize_pipeline()image_generator = ImageGeneratorAndSaver(pipe, self.config)image_generator.generate_and_save_images(data)if __name__ == "__main__":manager = ImageGenerationManager()manager.run()

整体流程

程序首先对配置参数进行初始化，接着加载提示信息，然后初始化模型管道，之后按照提示信息生成图像并保存，最后将生成的图像保存到指定文件夹。整个过程由 ImageGenerationManager 类进行协调管理。

1. 配置参数初始化

• 单例模式配置类 ConfigSingleton：
  • 借助 argparse 模块解析命令行参数，这些参数涵盖了模型路径、图像的高度和宽度、引导比例、推理步数、每个提示生成的图像数量、包含提示信息的 JSON 文件路径以及保存生成图像的文件夹路径。
  • 把解析后的参数存储在 config 字典中，同时创建保存图像的文件夹（若该文件夹不存在）。

2. 数据加载

• 数据加载器类 DataLoader：
  • 接收配置信息作为输入。
  • 尝试打开并读取 JSON 文件，若文件不存在或者解析失败，会输出相应的错误信息并返回空列表。

3. 模型管道初始化

• 模型管道初始化类 PipelineInitializer：
  • 接收配置信息作为输入。
  • 利用 CogView4Pipeline.from_pretrained 方法加载预训练模型，并将其数据类型设定为 torch.bfloat16。
  • 把模型移到 GPU 上（使用 pipe.to("cuda")），开启 vae 的切片和分块功能以减少 GPU 内存的使用。

4. 文件名生成

• 文件名生成器类 FilenameGenerator：
  • 接收配置信息作为输入。
  • 生成独一无二的文件名，文件名由当前时间戳和随机字符组合而成。

5. 图像生成与保存

• 图像生成和保存类 ImageGeneratorAndSaver：
  • 接收模型管道和配置信息作为输入。
  • 针对 JSON 文件中的每个提示信息：
    • 运用 torch.Generator("cuda") 创建一个 CUDA 随机数生成器。
    • 调用模型管道的 pipe 方法生成图像，传入提示信息、引导比例、推理步数、图像的高度和宽度等参数。
    • 为每张生成的图像生成一个唯一的文件名，并将其保存到指定文件夹，同时输出保存信息。
    • 若生成图像时出现异常，会输出相应的错误信息。

6. 主程序执行

• 主类 ImageGenerationManager：
  • 初始化配置信息、数据加载器和模型管道初始化器。
  • 调用 run 方法：
    • 加载 JSON 文件中的提示信息。
    • 若提示信息不为空，初始化模型管道。
    • 创建图像生成和保存类的实例，并调用其 generate_and_save_images 方法生成并保存图像。