CSGO: Content-Style Composition in Text-to-Image Generation（代码的复现）

CSGO简介

CSGO: Content-Style Composition in Text-to-Image Generation（风格迁移）
本文是一篇风格迁移的论文：将内容参考图像和风格参考图像分别投影，然后注入到内容模块和风格模块，同时采用controlnet的方法将内容参考图像注入unet的上采样块当中。

github中项目的地址

论文的代码部署

需要下载的模型权重：

我们的方法与 SDXL、VAE、ControlNet 和图像编码器完全兼容。请下载它们并将它们放在 ./base_models 文件夹中。
按照readme里面的指引，下载到如下文件夹里面：

复现中存在的一些问题

①需要保证如下包的版本与readme一致

diffusers==0.25.1
torch==2.0.1
torchaudio==2.0.2
torchvision==0.15.2
transformers==4.40.2

② NotImplementedError: Cannot copy out of meta tensor; no data!
参考知乎这篇
大语言模型调用踩坑点记录
数据在显存和内存中切换，导致出问题（显存不够）

部分参数从gpu拷贝到cpu会报错，将改成low_cpu_mem_usage=False，可以正常推理

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    base_model_path,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    add_watermarker=False,
    use_safetensors=True,
    vae=vae,
    revision="fp16",
	##这个参数
    low_cpu_mem_usage=False

)

③模型加载的问题
由于下载的模型权重都是fp16的格式的，然而这里模型的加载方式的参数是统一在最外面控制的，导致不同模型加载时，识别不了对应的模型文件：
加载模型是一些参数的设定

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    base_model_path,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    add_watermarker=False,
    #用safetensors格式的权重文件
    use_safetensors=True,
    vae=vae,
    revision="fp16",
    #device_map="auto"
    low_cpu_mem_usage=False

)

#这两个参数同时为fp16才会去读fp16的文件
 revision="fp16"
 variant= "fp16"

④需要统一数据的数据类型
由于之前的文本编码器的权重读取的是fp32，导致后续出现数据的类型不相同不能做运算的情况。
将pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
base_model_path,
controlnet=controlnet,
torch_dtype=torch.float16,
add_watermarker=False,
use_safetensors=True,
vae=vae,
revision=“fp16”,
#device_map=“auto”
low_cpu_mem_usage=False

)中的.from_pretrained函数（这个函数在pipeline_utils.py文件夹里）进行修改，当模型是文本编码器时，修改传入的一些参数

 if name == "text_encoder":
 					#如果是文本编码器，将varient设置为fp16
                    variant = "fp16"
                    if variant is not None:
                        # for folder in os.listdir(cached_folder):
                        folder_path = os.path.join(cached_folder, "text_encoder")
                        is_folder = os.path.isdir(folder_path) and "text_encoder" in config_dict
                        variant_exists = is_folder and any(
                            p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)
                        )
                        if variant_exists:
                            model_variants["text_encoder"] = variant
                if name == "text_encoder_2":
                    variant = "fp16"
                    if variant is not None:
                        # for folder in os.listdir(cached_folder):
                        folder_path = os.path.join(cached_folder, "text_encoder_2")
                        is_folder = os.path.isdir(folder_path) and "text_encoder_2" in config_dict
                        variant_exists = is_folder and any(
                            p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)
                        )
                        if variant_exists:
                            model_variants["text_encoder_2"] = variant
                            
                loaded_sub_model = load_sub_model(
                    library_name=library_name,
                    class_name=class_name,
                    importable_classes=importable_classes,
                    pipelines=pipelines,
                    is_pipeline_module=is_pipeline_module,
                    pipeline_class=pipeline_class,
                    torch_dtype=torch_dtype,
                    provider=provider,
                    sess_options=sess_options,
                    device_map=device_map,
                    max_memory=max_memory,
                    offload_folder=offload_folder,
                    offload_state_dict=offload_state_dict,
                    model_variants=model_variants,
                    name=name,
                    from_flax=from_flax,
                    variant=variant,
                    low_cpu_mem_usage=low_cpu_mem_usage,
                    cached_folder=cached_folder,
                    revision=revision,
                )
                logger.info(
                    f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}."
                )
               

⑤推理代码中的风格图像和内容图像都要是两个图像，而给的代码中是一个文本，一个图像

## 注意这里的两个图片都要转化为图片的格式，论文给的推理代码一个是文本，另一个是
style_image = Image.open("/mnt/CSGO-main/assets/{}".format(style_name)).convert('RGB')
content_image = Image.open('/mnt/test/image/{}'.format(content_name)).convert('RGB')

推理代码

import random

import torch
from ip_adapter.utils import BLOCKS as BLOCKS
from ip_adapter.utils import controlnet_BLOCKS as controlnet_BLOCKS
from PIL import Image
from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    StableDiffusionXLControlNetPipeline,

)
from ip_adapter import CSGO


#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

device = 'cuda:0'

base_model_path =  "./base_models/stable-diffusion-xl-base-1.0"  
image_encoder_path = "./base_models/IP-Adapter/sdxl_models/image_encoder"
csgo_ckpt = "./CSGO/csgo4_32.bin"
pretrained_vae_name_or_path ='./base_models/vae'
controlnet_path = "./base_models/TTPLanet_SDXL_Controlnet_Tile_Realistic"
weight_dtype = torch.float16
weight_dtype = torch.float16


vae = AutoencoderKL.from_pretrained(pretrained_vae_name_or_path,torch_dtype=torch.float16)
controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16,use_safetensors=True)

def get_device_map():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

device = get_device_map()

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    base_model_path,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    add_watermarker=False,
    use_safetensors=True,
    vae=vae,
    revision="fp16",
    ## 这里要加这个代码，不然会报错，因为显存不够，然后导致数据在显存和内存之间转换，报错
    low_cpu_mem_usage=False
)
pipe.enable_vae_tiling()

target_content_blocks = BLOCKS['content']
target_style_blocks = BLOCKS['style']
controlnet_target_content_blocks = controlnet_BLOCKS['content']
controlnet_target_style_blocks = controlnet_BLOCKS['style']


csgo = CSGO(pipe, image_encoder_path, csgo_ckpt, device, num_content_tokens=4,num_style_tokens=32,
                          target_content_blocks=target_content_blocks, target_style_blocks=target_style_blocks,controlnet_adapter=True,
                              controlnet_target_content_blocks=controlnet_target_content_blocks,
                              controlnet_target_style_blocks=controlnet_target_style_blocks,
                              content_model_resampler=True,
                              style_model_resampler=True,


                              )

style_name = 'img_0.png'
content_name = 's_01_e_26_shot_005126_005200.png'

## 注意这里的两个图片都要转化为图片的格式，论文给的推理代码一个是文本，另一个是
style_image = Image.open("/mnt/CSGO-main/assets/{}".format(style_name)).convert('RGB')
content_image = Image.open('/mnt/test/image/{}'.format(content_name)).convert('RGB')

num_sample=1
caption = ''
#写个循环，看看各个参数对生成图片的影响
while True:
    tem = 0
    for ccs in range(5, 11, 1):
        ccs = ccs * 0.1
        content_scale = random.uniform(0.6, 1.5)
        style_scale = random.uniform(0.5, 1)
        images = csgo.generate(pil_content_image= content_image, pil_style_image=style_image,
                                   prompt=caption,
                                   negative_prompt= "text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry",
                                   content_scale=1.0,
                                   style_scale=1.0,
                                   guidance_scale=10,
                                   num_images_per_prompt=num_sample,
                                   num_samples=1,
                                   num_inference_steps=50,
                                   seed=42,
                                   image=content_image.convert('RGB'),
                                   controlnet_conditioning_scale=0.6,
                                  )
        formatted_ccs = "{:.2f}".format(ccs)
        formatted_content_scale = "{:.2f}".format(content_scale)
        formatted_style_scale = "{:.2f}".format(style_scale)
        images[0].save(f"inference/ccs:{formatted_ccs}-cs:{formatted_content_scale}-ss:{formatted_style_scale}.png")

    tem = tem + 1
    if tem >= 100:
        break

生成结果示意图

风格参考图像

文本： a cat

生成的图像