generate.py

from diffusers import DiffusionPipeline
import torch
from PIL import Image


def test_lora(lcm_speedup=False):
    pipe = DiffusionPipeline.from_pretrained("/mnt/workspace/Project/VideoGen/ckpt/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker = None, requires_safety_checker=False)
    pipe.to("cuda")

    lora_path = "/mnt/workspace/Project/VideoGen/TrainT2V/diffusers/examples/text_to_image/results" 
    pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="artifacts_100epoch_lora.safetensors", adapter_name="pattern")
    if lcm_speedup:
        pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="lcm_lora.safetensors", adapter_name="lcm")
        pipe.set_adapters(["artifacts", "lcm"], adapter_weights=[1.0, 1.0])

    prompts = [
        # "Elegant 19th-century crystal glass collection, exuding vintage charm with intricate fluted designs and elegant contours. Each piece shines with timeless craftsmanship, featuring subtle patterns and sturdy bases. This exquisite set reflects the opulence of Victorian-era glasswork, used in fine dining to add sophistication and allure to dining experiences. The translucent quality and ornate details highlight a bygone era's artisanal skill, preserved in immaculate condition.\"\n\n\"Elegant 19th-century crystal glass collection, exuding vintage charm with intricate fluted designs and elegant contours. Each piece shines with timeless craftsmanship, featuring subtle patterns and sturdy bases. This exquisite set reflects the opulence of Victorian-era glasswork, used in fine dining to add sophistication and allure to dining experiences. The translucent quality and ornate details highlight a bygone era's artisanal skill, preserved in immaculate condition.",
        # "A vintage camera with a rectangular body, equipped with two dials and ornate decorative patterns. Its boxy shape and weathered surface hint at early 20th-century craftsmanship. Historically significant for documenting pioneer moments, this antique camera is a testament to early visual storytelling. Its intricate designs highlight the meticulous work of artisans of the time.",
        "A captivating intricately patterned ceramic vessel illuminated by a soft inner light, showcasing an elegant array of flowing curves and organic shapes. Crafted from delicate porcelain, its minimalist yet complex design evokes a sense of serenity. The interplay of light and shadow accentuates its textures and fine lines, symbolizing traditional artistry merged with modern aesthetics. This artifact, although not explicitly marked with age, exudes a timeless elegance reflective of master craftsmanship.",
        "Ancient Bronze Sword Guard: This intricately carved sword guard is a masterpiece of Chinese craftsmanship from the Western Zhou period, dating back to around 900 BCE. The green patina of oxidized bronze is a testament to its age. The guard is adorned with delicate spiral and interlocking patterns, symbolizing both beauty and power. The ornate design not only protects the blade but also showcases the artistry of its time, reflecting the society's high regard for elegance and strength. Today, it stands as a historical relic, silently narrating tales of a bygone era."
    ]

    # 设置生成参数
    if lcm_speedup:
        num_inference_steps = 8
        guidance_scale = 2
    else:
        num_inference_steps = 30
        guidance_scale = 7.5
        
    num_samples_per_prompt = 3

    # 创建一个空的图像列表
    all_images = []

    # 为每个 prompt 生成 num_samples_per_prompt 张图片
    for prompt in prompts:
        images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_samples_per_prompt).images
        all_images.extend(images)

    # 创建一个 2x3 的网格图
    grid_image = Image.new('RGB', (3 * 512, 2 * 512))  # 假设每张图片大小为 512x512
    for idx, img in enumerate(all_images):
        x = (idx % 3) * 512
        y = (idx // 3) * 512
        grid_image.paste(img, (x, y))

    # 保存网格图
    n = 4 if lcm_speedup else 30
    grid_image.save(f"test_lora_grid_{n}_steps.png")


if __name__ == "__main__":
    test_lora()