Datawhale X 魔搭 AI夏令营夏令营第四期 Task2

Sugarhater

已于 2024-08-14 23:56:34 修改

阅读量307

点赞数 9

文章标签：人工智能深度学习

于 2024-08-14 23:42:58 首次发布

本文链接：https://blog.csdn.net/Sugarhater/article/details/141202816

版权

AI生图产生的问题

难以区别图片真假，产生DeepFake问题。解决这一问题，可以使用CNN进行识别。

代码解读

下载必要的包

!pip install simple-aesthetics-predictor
!pip install -v -e data-juicer
!pip uninstall pytorch-lightning -y
!pip install peft lightning pandas torchvision
!pip install -e DiffSynth-Studio

下载、训练和加载训练好的模型

from diffsynth import download_models
download_models(["Kolors", "SDXL-vae-fp16-fix"])

import os

cmd = """
python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py \
  --pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors \
  --pretrained_text_encoder_path models/kolors/Kolors/text_encoder \
  --pretrained_fp16_vae_path models/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors \
  --lora_rank 16 \
  --lora_alpha 4.0 \
  --dataset_path data/lora_dataset_processed \
  --output_path ./models \
  --max_epochs 1 \
  --center_crop \
  --use_gradient_checkpointing \
  --precision "16-mixed"
""".strip()

os.system(cmd)

from diffsynth import ModelManager, SDXLImagePipeline
from peft import LoraConfig, inject_adapter_in_model
import torch

# 目的是加载一个 LoRA 模型，并将其注入到预训练模型中
def load_lora(model, lora_rank, lora_alpha, lora_path):
    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=["to_q", "to_k", "to_v", "to_out"],
    )
    model = inject_adapter_in_model(lora_config, model)
    state_dict = torch.load(lora_path, map_location="cpu")
    model.load_state_dict(state_dict, strict=False)
    return model


# 用于管理模型的加载，包括指定数据类型（torch.float16）和设备（cuda），并加载模型的权重文件。模型的文件路径通过 file_path_list 参数提供
model_manager = ModelManager(torch_dtype=torch.float16, device="cuda",
                             file_path_list=[
                                 "models/kolors/Kolors/text_encoder",
                                 "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors",
                                 "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors"
                             ])
pipe = SDXLImagePipeline.from_model_manager(model_manager)
# 将训练好的 LoRA 模型加载到 pipe.unet（图像生成管道中的 U-Net 模块）中。U-Net 是扩散模型的核心部分，负责图像的逐步生成和细化
pipe.unet = load_lora(
    pipe.unet,
    lora_rank=16, # 
    lora_alpha=2.0
    lora_path="models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt"
)

调整提示词，生成图像

cfg_scale：表示 Classifier-Free Guidance Scale，用于控制生成图像时的提示强度。值越高，生成的图像越会严格遵循提示，但过高的值可能导致图像质量下降或出现不自然的效果。
num_inference_steps：这是推理步骤的数量。模型在生成图像时的迭代次数，更多的迭代通常会带来更高质量的图像，但也会增加生成时间。
height=1024, width=1024：生成图像的高度和宽度，单位为像素。这表示生成的图像将是 1024x1024 像素的正方形图像。