AIGC baseline1 学习

4467 字
12 分钟

baseline 分析

首先总体分析baseline代码,这段代码可以分为以下几个部分: 1. 安装依赖:使用pip命令安装所需的Python库,如simple-aesthetics-predictor、data-juicer、pandas等,并卸载旧版本的pytorch-lightning。 2. 数据加载与预处理:从MsDataset加载数据集,保存图像并创建带有元数据的JSONL文件。配置data-juicer来筛选图像数据。 3. 模型推理与数据增强:使用CLIPModel进行图像-文本相似度计算,并通过StableDiffusionPipeline生成二次元图像。 4. 图像生成与保存:生成多幅二次元图像并保存,通过将这些图像拼接成一个更大的图像进行展示。

这些步骤构成了从数据加载、预处理、模型推理到最终图像生成的完整工作流程。

接下来,我们对于这几个部分的代码逐步分析,拆解来看

数据加载与预处理

  1. 数据集加载:
from modelscope.msdatasets import MsDataset

ds = MsDataset.load(
    'AI-ModelScope/lowres_anime',
    subset_name='default',
    split='train',
    cache_dir="/mnt/workspace/kolors/data"
)

这里使用MsDataset从ModelScope加载了一个低分辨率的动漫数据集,并将其存储在指定的缓存目录中。 2. 图像与元数据处理:

import json, os
from tqdm import tqdm
os.makedirs("./data/lora_dataset/train", exist_ok=True)
os.makedirs("./data/data-juicer/input", exist_ok=True)
with open("./data/data-juicer/input/metadata.jsonl", "w") as f:
    for data_id, data in enumerate(tqdm(ds)):
        image = data["image"].convert("RGB")
        image.save(f"/mnt/workspace/kolors/data/lora_dataset/train/{data_id}.jpg")
        metadata = {"text": "二次元", "image": [f"/mnt/workspace/kolors/data/lora_dataset/train/{data_id}.jpg"]}
        f.write(json.dumps(metadata))
        f.write("\n")
* 文件夹创建:代码创建了两个文件夹train和input,分别用于保存图像和元数据。
* 数据处理:遍历数据集中的每个样本,将图像转换为RGB格式,并保存到指定目录。随后,将每个图像的路径与对应的文本描述(这里是固定的“二次元”)打包成一个JSON对象,并写入metadata.jsonl文件。

3. 配置文件生成:

data_juicer_config = """
#global parameters
project_name: 'data-process'
dataset_path: './data/data-juicer/input/metadata.jsonl'  # path to your dataset directory or file
np: 4  # number of subprocess to process your dataset

text_keys: 'text'
image_key: 'image'
image_special_token: '<__dj__image>'

export_path: './data/data-juicer/output/result.jsonl'

#process schedule
#a list of several process operators with their arguments
process:
    - image_shape_filter:
        min_width: 1024
        min_height: 1024
        any_or_all: any
    - image_aspect_ratio_filter:
        min_ratio: 0.5
        max_ratio: 2.0
        any_or_all: any
"""
with open("data/data-juicer/data_juicer_config.yaml", "w") as file:
    file.write(data_juicer_config.strip())
* 配置文件生成:生成data_juicer的配置文件,用于数据过滤和处理。配置文件指定了项目名称、数据集路径、文本和图像键值、输出路径以及图像筛选的条件。

4. 数据处理执行:

!dj-process --config data/data-juicer/data_juicer_config.yaml

使用生成的配置文件执行data-juicer的数据处理流程,最终将符合条件的图像和对应的文本保存在指定的输出文件中。

这部分代码主要完成了数据集的加载、图像预处理、元数据生成以及数据筛选的配置和执行,为后续的模型训练和推理打下了基础。

模型推理与数据增强

  1. 加载CLIP模型和处理器:
from transformers import CLIPProcessor, CLIPModel
import torch

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
* 加载预训练的CLIPModel和CLIPProcessor用于处理图像和文本数据。

2. 计算图像-文本相似度:

image = [Image.open(img_path) for img_path in df["file_name"]]
inputs = processor(text=df["text"].tolist(), image=image, return_tensors="pt", padding=True)

outputs = model(**inputs)
logits_per_image = outputs.logits_per_image
probs = logits_per_image.softmax(dim=1)
* 将文本和图像数据输入到模型中,计算它们之间的相似度得分,并通过softmax获得概率分布。

3. 数据集和数据加载器:

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, df, processor):
        self.texts = df["text"].tolist()
        self.image = [Image.open(img_path) for img_path in df["file_name"]]
        self.processor = processor

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        inputs = self.processor(text=self.texts[idx], image=self.image[idx], return_tensors="pt", padding=True)
        return inputs

dataset = CustomDataset(df, processor)
dataloader = DataLoader(dataset, batch_size=8)
* 定义自定义数据集类CustomDataset,将文本和图像处理成模型的输入格式,并使用DataLoader批量处理数据。

4. 批量推理:

for batch in dataloader:
    outputs = model(**batch)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)
    print(probs)
* 对批量数据进行推理,并打印图像与文本之间的相似度概率。

图像生成与保存

  1. 加载Stable Diffusion管道:
from diffusers import StableDiffusionPipeline

torch.manual_seed(1)
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
* 加载预训练的StableDiffusionPipeline模型,并将其部署到GPU上以提高推理速度。

2. 生成图像:

prompt = "二次元,一个紫色长发小女孩穿着粉色吊带漏肩连衣裙,在练习室练习唱歌,手持话筒"
negative_prompt = "丑陋、变形、嘈杂、模糊、低对比度"
guidance_scale = 4
num_inference_steps = 50

image = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    guidance_scale=guidance_scale,
    num_inference_steps=num_inference_steps,
    height=1024,
    width=1024,
).image[0]

image.save("example_image.png")
image
* 使用给定的提示词prompt和negative_prompt生成图像,并将图像保存为文件。

3. 批量生成不同场景的图像:

torch.manual_seed(1)
image = pipe(
    prompt="二次元,日系动漫,演唱会的观众席,人山人海,一个紫色短发小女孩穿着粉色吊带漏肩连衣裙坐在演唱会的观众席,舞台上衣着华丽的歌星们在唱歌",
    negative_prompt="丑陋、变形、嘈杂、模糊、低对比度",
    cfg_scale=4,
    num_inference_steps=50, height=1024, width=1024,
)
image.save("1.jpg")
* 使用不同的提示词生成一系列图像,并保存为多个文件。

4. 拼接图像:

import numpy as np
from PIL import Image

image = [np.array(Image.open(f"{i}.jpg")) for i in range(1, 9)]
image = np.concatenate([
    np.concatenate(image[0:2], axis=1),
    np.concatenate(image[2:4], axis=1),
    np.concatenate(image[4:6], axis=1),
    np.concatenate(image[6:8], axis=1),
], axis=0)
image = Image.fromarray(image).resize((1024, 2048))
image
* 将生成的图像按照指定顺序拼接成一个大图,并调整尺寸用于展示。