run_txt2img_axe_infer_new.py

from typing import List, Union
import numpy as np
# import onnxruntime
import axengine
import torch
from PIL import Image
from transformers import CLIPTokenizer, CLIPTextModel, PreTrainedTokenizer, CLIPTextModelWithProjection

import time
import argparse
import uuid  # 用于生成唯一文件名
import os


def get_args():
    parser = argparse.ArgumentParser(
        prog="StableDiffusion",
        description="Generate picture with the input prompt"
    )
    parser.add_argument("--prompt", type=str, required=False, default="Self-portrait oil painting, a beautiful cyborg with golden hair, 8k", help="the input text prompt")
    parser.add_argument("--text_model_dir", type=str, required=False, default="./models/", help="Path to text encoder and tokenizer files")
    parser.add_argument("--unet_model", type=str, required=False, default="./models/unet.axmodel", help="Path to unet axmodel model")
    parser.add_argument("--vae_decoder_model", type=str, required=False, default="./models/vae_decoder.axmodel", help="Path to vae decoder axmodel model")
    parser.add_argument("--time_input", type=str, required=False, default="./models/time_input_txt2img.npy", help="Path to time input file")
    parser.add_argument("--save_dir", type=str, required=False, default="./txt2img_output_axe", help="Path to the output image file")
    return parser.parse_args()

def maybe_convert_prompt(prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"):  # noqa: F821
    if not isinstance(prompt, List):
        prompts = [prompt]
    else:
        prompts = prompt

    prompts = [_maybe_convert_prompt(p, tokenizer) for p in prompts]

    if not isinstance(prompt, List):
        return prompts[0]

    return prompts


def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noqa: F821
    tokens = tokenizer.tokenize(prompt)
    unique_tokens = set(tokens)
    for token in unique_tokens:
        if token in tokenizer.added_tokens_encoder:
            replacement = token
            i = 1
            while f"{token}_{i}" in tokenizer.added_tokens_encoder:
                replacement += f" {token}_{i}"
                i += 1

            prompt = prompt.replace(token, replacement)

    return prompt


def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
    tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)

    text_inputs = tokenizer(
        prompt,
        padding="max_length",
        max_length=77,
        truncation=True,
        return_tensors="pt",
    )
    text_input_ids = text_inputs.input_ids

    text_encoder = axengine.InferenceSession(
        os.path.join(
            text_encoder_dir,
            "sd15_text_encoder_sim.axmodel"
        ),
    )
    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy().astype(np.int32)})[0]

    prompt_embeds_npy = text_encoder_onnx_out
    return prompt_embeds_npy


def get_alphas_cumprod():
    betas = torch.linspace(0.00085 ** 0.5, 0.012 ** 0.5, 1000, dtype=torch.float32) ** 2
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0).detach().numpy()
    final_alphas_cumprod = alphas_cumprod[0]
    self_timesteps = np.arange(0, 1000)[::-1].copy().astype(np.int64)
    return alphas_cumprod, final_alphas_cumprod, self_timesteps



if __name__ == '__main__':
    args = get_args()
    
    tokenizer_dir = args.text_model_dir + 'tokenizer'
    text_encoder_dir = args.text_model_dir + 'text_encoder'
    unet_model = args.unet_model
    vae_decoder_model = args.vae_decoder_model
    time_input = args.time_input
    save_dir = args.save_dir

    # 确保保存目录存在
    os.makedirs(save_dir, exist_ok=True)

    print(f"tokenizer: {tokenizer_dir}")
    print(f"text_encoder: {text_encoder_dir}")
    print(f"unet_model: {unet_model}")
    print(f"vae_decoder_model: {vae_decoder_model}")
    print(f"time_input: {time_input}")
    print(f"save_dir: {save_dir}")

    # 加载模型（只加载一次）
    start = time.time()
    unet_session_main = axengine.InferenceSession(unet_model)
    vae_decoder = axengine.InferenceSession(vae_decoder_model)
    print(f"load models take {(1000 * (time.time() - start)):.1f}ms")

    # 主循环：支持多次输入 Prompt
    while True:
        # 用户输入 Prompt
        prompt = input("\nEnter a prompt to generate an image (or type 'exit' to quit): ")
        if prompt.lower() == 'exit':
            print("Exiting...")
            break

        # Text Encoder
        start = time.time()
        prompt_embeds_npy = get_embeds(prompt, tokenizer_dir, text_encoder_dir)
        print(f"text encoder take {(1000 * (time.time() - start)):.1f}ms")

        # 初始化 Latent
        latents_shape = [1, 4, 64, 64]
        latent = torch.randn(latents_shape, generator=None, device="cpu", dtype=torch.float32,
                             layout=torch.strided).detach().numpy()
                             
        alphas_cumprod, final_alphas_cumprod, self_timesteps = get_alphas_cumprod()

        # 加载 time_input 文件
        time_input_data = np.load(time_input)

        # U-Net Inference Loop
        timesteps = np.array([999, 759, 499, 259]).astype(np.int64)
        unet_loop_start = time.time()
        for i, timestep in enumerate(timesteps):
            unet_start = time.time()
            noise_pred = unet_session_main.run(None, {
                "sample": latent,
                "/down_blocks.0/resnets.0/act_1/Mul_output_0": np.expand_dims(time_input_data[i], axis=0),
                "encoder_hidden_states": prompt_embeds_npy
            })[0]
            print(f"unet once take {(1000 * (time.time() - unet_start)):.1f}ms")

            sample = latent
            model_output = noise_pred
            if i < 3:
                prev_timestep = timesteps[i + 1]
            else:
                prev_timestep = timestep

            alpha_prod_t = alphas_cumprod[timestep]
            alpha_prod_t_prev = alphas_cumprod[prev_timestep] if prev_timestep >= 0 else final_alphas_cumprod

            beta_prod_t = 1 - alpha_prod_t
            beta_prod_t_prev = 1 - alpha_prod_t_prev

            scaled_timestep = timestep * 10
            c_skip = 0.5 ** 2 / (scaled_timestep ** 2 + 0.5 ** 2)
            c_out = scaled_timestep / (scaled_timestep ** 2 + 0.5 ** 2) ** 0.5
            predicted_original_sample = (sample - (beta_prod_t ** 0.5) * model_output) / (alpha_prod_t ** 0.5)

            denoised = c_out * predicted_original_sample + c_skip * sample

            if i != 3:
                noise = torch.randn(model_output.shape, generator=None, device="cpu", dtype=torch.float32,
                                    layout=torch.strided).to("cpu").detach().numpy()
                prev_sample = (alpha_prod_t_prev ** 0.5) * denoised + (beta_prod_t_prev ** 0.5) * noise
            else:
                prev_sample = denoised

            latent = prev_sample

        print(f"unet loop take {(1000 * (time.time() - unet_loop_start)):.1f}ms")

        # VAE Inference
        vae_start = time.time()
        latent = latent / 0.18215
        image = vae_decoder.run(None, {"x": latent})[0]
        print(f"vae inference take {(1000 * (time.time() - vae_start)):.1f}ms")

        # 保存结果
        save_start = time.time()
        image = np.transpose(image, (0, 2, 3, 1)).squeeze(axis=0)
        image_denorm = np.clip(image / 2 + 0.5, 0, 1)
        image = (image_denorm * 255).round().astype("uint8")
        pil_image = Image.fromarray(image[:, :, :3])

        # 使用 UUID 确保文件名唯一
        unique_filename = f"{uuid.uuid4()}.png"
        save_path = os.path.join(save_dir, unique_filename)
        pil_image.save(save_path)
        print(f"Image saved to {save_path}")
        print(f"Save image take {(1000 * (time.time() - save_start)):.1f}ms")