bghira / SimpleTuner

A general fine-tuning kit geared toward diffusion models.
GNU Affero General Public License v3.0
1.79k stars 171 forks source link

kolors inference code? #1169

Open Deng-Xian-Sheng opened 2 hours ago

Deng-Xian-Sheng commented 2 hours ago

I not find kolors inference code.

I not use inference , if no code or doc.

I use ChatGPT write this is code, work, but exist no exist very good code?

import torch
from diffusers import KolorsPipeline
from safetensors.torch import load_file
import os
from datetime import datetime
# 1. 加载基础模型
pipe = KolorsPipeline.from_pretrained(
    "Kwai-Kolors/Kolors-diffusers", 
    torch_dtype=torch.float16, 
    variant="fp16"
).to("cuda")
# 2. 加载训练好的权重
checkpoint_path = "checkpoint-27304/unet/diffusion_pytorch_model.safetensors"
state_dict = load_file(checkpoint_path)
pipe.unet.load_state_dict(state_dict)
pipe.unet.eval()
# 3. 创建输出目录
output_dir = "outputs"
os.makedirs(output_dir, existok=True)
# 4. 生成和保存图片
prompt = '一个女孩穿着粉红色的裙子,手里捧着鲜花,旁边还有一个小男孩,右下角有个水印叫“@邓文怡”'
seed = 20091114
# 生成时间戳
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
# 构建文件名(包含主要信息)
filename = f"{timestamp}_seed{seed}.png"
# 或者更详细的文件名:
# filename = f"{timestamp}_seed{seed}_steps{50}_scale{7.5}.png"
output_path = os.path.join(output_dir, filename)
# 生成图片
image = pipe(
    prompt=prompt,
    negative_prompt="昏暗、扭曲、模糊、不真实、nfsw",
    guidance_scale=5.0,
    num_inference_steps=50,
    generator=torch.Generator(pipe.device).manual_seed(seed),
).images[0]
# 保存图片
image.save(output_path)
print(f"Image saved to: {output_path}")
bghira commented 2 hours ago

the inference code is the same as for SDXL

Deng-Xian-Sheng commented 2 hours ago

ok thank you help!!