Upload run_txt2img_axe_infer_new.py

Only load the model once, support multiple rounds of Prompt to generate images and save to the specified folder

run_txt2img_axe_infer_new.py

@@ -0,0 +1,197 @@
+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_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
+                                                 torch_dtype=torch.float32,
+                                                 variant="fp16")
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=77,
+        truncation=True,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
+
+    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    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")