add ax620e_320x320_models & gradio_demo.py

ax620e_320x320_models/text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/home/patrick/.cache/huggingface/hub/models--lykon-models--dreamshaper-7/snapshots/c4c9f9bec821e1862a78cbf45685cfb35b93638d/text_encoder",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float16",
+  "transformers_version": "4.33.0.dev0",
+  "vocab_size": 49408
+}

ax620e_320x320_models/text_encoder/sd15_text_encoder_sim.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d348ba3a0f0c70552b92215a8f78496f1c2072364e393510e0101af382fbcf4
+size 240153225

ax620e_320x320_models/time_input_img2img.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95d015256308e1be1af00793c77fa2ba8c934163beaa8015dec54d20048838cf
+size 20608

ax620e_320x320_models/time_input_txt2img.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a48a430879c6a81a907889a6bb2b73f48cc9dc45b52f047ad2ee5c2dddcd2d10
+size 20608

ax620e_320x320_models/tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

ax620e_320x320_models/tokenizer/tokenizer_config.json

@@ -0,0 +1,33 @@
+{
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

ax620e_320x320_models/unet.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb7c119d421cac390ffc4c03305d865d020cacb3aed528a4b06bf8007dfaf78d
+size 969190063

ax620e_320x320_models/vae_decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36c5f95c7ea5f94cdd5d76814e6b098fcff8e78c82faecd1e78723429f275e23
+size 94370744

ax620e_320x320_models/vae_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4bd0815ff6220d79e12aaaa526a269280e86918020e6fc576050163984f57a4
+size 198057245

ax620e_320x320_models/vae_encoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9124d11155d4ff6470a416264bae5412240e1ebe8c2b53e79935f8bcdfde0b8
+size 60221332

ax620e_320x320_models/vae_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:386ca46adcf4b612b1ffa8cff88dd06e5c53c6292e6da161b27c088734815d03
+size 136728111

gradio_demo.py

@@ -0,0 +1,420 @@
+from typing import Optional, Tuple, Union
+from functools import lru_cache
+import argparse
+import os
+import warnings
+import socket
+
+import numpy as np
+import torch
+import gradio as gr
+
+REQUIRED_GRADIO_VERSION = "5.42.0"
+from PIL import Image
+
+from diffusers.utils import make_image_grid, load_image
+from diffusers.utils.torch_utils import randn_tensor
+
+from launcher import (
+    resolve_dimensions,
+    compute_latent_shape,
+    get_embeds,
+    get_alphas_cumprod,
+    create_session,
+    prepare_init_image,
+    ensure_parent,
+    resolve_with_base,
+    add_noise,
+    retrieve_latents,
+    denoise_loop,
+    DiagonalGaussianDistribution,
+    AutoencoderKLOutput,
+    IMG2IMG_TIMESTEPS,
+    IMG2IMG_SELF_TIMESTEPS,
+    IMG2IMG_STEP_INDEX,
+    TXT2IMG_TIMESTEPS,
+    TIME_EMBED_KEY,
+)
+
+
+@lru_cache(maxsize=8)
+def _cached_session(model_path: str, backend: str):
+    return create_session(model_path, backend)
+
+
+def _check_gradio_version():
+    import gradio
+    ver = getattr(gradio, "__version__", None)
+    if ver is None or ver.split("+")[0] != REQUIRED_GRADIO_VERSION:
+        warnings.warn(
+            f"当前 gradio 版本为 {ver}, 建议使用 {REQUIRED_GRADIO_VERSION} 以避免兼容性问题。"
+        )
+
+
+def _preload_models(model_dir: str, backend: str, isize: Union[int, str]):
+    """在前端启动前加载关键模型/输入，提前暴露加载错误。"""
+    backend = backend.lower()
+    model_suffix = ".axmodel" if backend == "axe" else ".onnx"
+    text_encoder_path = os.path.join(model_dir, "text_encoder", f"sd15_text_encoder_sim{model_suffix}")
+    unet_model = os.path.join(model_dir, f"unet{model_suffix}")
+    vae_decoder_model = os.path.join(model_dir, f"vae_decoder{model_suffix}")
+    vae_encoder_model = os.path.join(model_dir, f"vae_encoder{model_suffix}")
+
+    # 提前加载会话
+    _cached_session(text_encoder_path, backend)
+    _cached_session(unet_model, backend)
+    _cached_session(vae_decoder_model, backend)
+    # vae encoder 仅 img2img 用，若不存在可跳过
+    if os.path.exists(vae_encoder_model):
+        _cached_session(vae_encoder_model, backend)
+
+    # 预加载时间输入文件，txt2img & img2img
+    txt2img_time = os.path.join(model_dir, "time_input_txt2img.npy")
+    img2img_time = os.path.join(model_dir, "time_input_img2img.npy")
+    for path in (txt2img_time, img2img_time):
+        if os.path.exists(path):
+            np.load(path)
+        else:
+            warnings.warn(f"缺少时间输入文件: {path}")
+
+    # 确认分辨率合法，主要为了提前暴露 isize 参数错误
+    resolve_dimensions(isize, None, None)
+
+
+def _list_host_ips() -> list:
+    ips = set()
+    try:
+        hostname = socket.gethostname()
+        infos = socket.getaddrinfo(hostname, None, family=socket.AF_INET)
+        for info in infos:
+            ip = info[4][0]
+            if ip and not ip.startswith("127."):
+                ips.add(ip)
+    except Exception:
+        pass
+    if not ips:
+        ips.add("127.0.0.1")
+    return sorted(ips)
+
+
+def _prepare_init_image_any(image_source: Union[str, Image.Image], height: int, width: int) -> Tuple[Image.Image, np.ndarray]:
+    if isinstance(image_source, Image.Image):
+        image = image_source.resize((width, height)).convert("RGB")
+        image_show = image.copy()
+        np_img = (np.array(image).astype(np.float32) / 255.0)[None, ...]
+        np_img = torch.from_numpy(np_img.transpose(0, 3, 1, 2)).numpy()
+        np_img = 2.0 * np_img - 1.0
+        return image_show, np_img
+    image_show, processed = prepare_init_image(str(image_source), height, width)
+    return image_show, processed
+
+
+def _denoise_loop(latent: np.ndarray,
+                  prompt_embeds: np.ndarray,
+                  time_inputs: np.ndarray,
+                  timesteps: np.ndarray,
+                  unet_session,
+                  alphas_cumprod: np.ndarray,
+                  final_alphas_cumprod: float,
+                  generator: Optional[torch.Generator],
+                  noise_dtype: torch.dtype,
+                  self_timesteps: Optional[np.ndarray] = None,
+                  step_index: Optional[list] = None) -> np.ndarray:
+    if time_inputs.shape[0] < len(timesteps):
+        raise ValueError("time_input 的步数少于推理步数")
+
+    device = torch.device("cpu")
+    for i, timestep in enumerate(timesteps):
+        latent = latent.astype(np.float32)
+        feeds = {
+            "sample": latent,
+            TIME_EMBED_KEY: np.expand_dims(time_inputs[i], axis=0),
+            "encoder_hidden_states": prompt_embeds,
+        }
+        noise_pred = unet_session.run(None, feeds)[0]
+
+        sample = latent
+        model_output = noise_pred
+        if self_timesteps is not None and step_index is not None:
+            prev_idx = step_index[i] + 1
+            if prev_idx < len(self_timesteps):
+                prev_timestep = int(self_timesteps[prev_idx])
+            else:
+                prev_timestep = int(timestep)
+        elif i + 1 < len(timesteps):
+            prev_timestep = int(timesteps[i + 1])
+        else:
+            prev_timestep = int(timestep)
+
+        alpha_prod_t = alphas_cumprod[int(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 = int(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 != len(timesteps) - 1:
+            if noise_dtype == torch.float32 and generator is None:
+                noise = torch.randn(model_output.shape, device=device, dtype=noise_dtype).cpu().numpy()
+            else:
+                noise_tensor = randn_tensor(model_output.shape, generator=generator, device=device, dtype=noise_dtype)
+                noise = noise_tensor.cpu().numpy()
+            prev_sample = (alpha_prod_t_prev ** 0.5) * denoised + (beta_prod_t_prev ** 0.5) * noise
+        else:
+            prev_sample = denoised
+
+        latent = prev_sample.astype(np.float32)
+
+    return latent
+
+
+def run_pipeline(prompt: str,
+                 model_dir: str = "./models",
+                 backend: str = "axe",
+                 isize: Union[int, str] = "512",
+                 height: Optional[int] = None,
+                 width: Optional[int] = None,
+                 seed: Optional[int] = None,
+                 time_input_override: Optional[str] = None,
+                 init_image: Optional[Union[str, Image.Image]] = None,
+                 save_path: Optional[str] = None):
+    backend = backend.lower()
+    is_img2img = init_image is not None
+
+    tokenizer_dir = os.path.join(model_dir, "tokenizer")
+    text_encoder_dir = os.path.join(model_dir, "text_encoder")
+
+    model_suffix = ".axmodel" if backend == "axe" else ".onnx"
+    text_encoder_path = os.path.join(text_encoder_dir, f"sd15_text_encoder_sim{model_suffix}")
+    unet_model = os.path.join(model_dir, f"unet{model_suffix}")
+    vae_decoder_model = os.path.join(model_dir, f"vae_decoder{model_suffix}")
+    vae_encoder_model = os.path.join(model_dir, f"vae_encoder{model_suffix}")
+    time_input_default = "time_input_img2img.npy" if is_img2img else "time_input_txt2img.npy"
+    if time_input_override:
+        time_input_path = resolve_with_base(time_input_override, model_dir)
+    else:
+        time_input_path = os.path.join(model_dir, time_input_default)
+
+    if isinstance(init_image, str):
+        init_image_source = resolve_with_base(init_image, model_dir)
+    else:
+        init_image_source = init_image
+
+    height, width = resolve_dimensions(isize, height, width)
+
+    device = torch.device("cpu")
+    if seed is None or int(seed) < 0:
+        seed_used = int(torch.seed())
+    else:
+        seed_used = int(seed)
+    generator: Optional[torch.Generator] = torch.manual_seed(seed_used)
+    noise_dtype = torch.float16 if is_img2img else torch.float32
+
+    prompt_embeds_npy = get_embeds(prompt, tokenizer_dir, text_encoder_path, backend)
+    alphas_cumprod, final_alphas_cumprod, _ = get_alphas_cumprod()
+
+    vae_encoder_session = _cached_session(vae_encoder_model, backend) if is_img2img else None
+    unet_session = _cached_session(unet_model, backend)
+    vae_decoder_session = _cached_session(vae_decoder_model, backend)
+
+    time_input = np.load(time_input_path)
+
+    if is_img2img:
+        init_image_show, init_image_np = _prepare_init_image_any(init_image_source, height, width)
+
+        vae_encoder_inp_name = vae_encoder_session.get_inputs()[0].name
+        vae_encoder_out = vae_encoder_session.run(None, {vae_encoder_inp_name: init_image_np})[0]
+
+        posterior = DiagonalGaussianDistribution(torch.from_numpy(vae_encoder_out).to(torch.float32))
+        vae_encode_info = AutoencoderKLOutput(latent_dist=posterior)
+        init_latents = retrieve_latents(vae_encode_info, generator=generator)
+        init_latents = init_latents * 0.18215
+        init_latents = torch.cat([init_latents], dim=0)
+        noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=noise_dtype)
+        timestep_tensor = torch.tensor([int(IMG2IMG_TIMESTEPS[0])], device=device)
+        init_latents = add_noise(init_latents.to(device), noise, timestep_tensor)
+        latent = init_latents.detach().cpu().numpy()
+
+        timesteps = IMG2IMG_TIMESTEPS
+        self_timesteps = IMG2IMG_SELF_TIMESTEPS
+        step_index = IMG2IMG_STEP_INDEX
+    else:
+        batch, channels, latent_h, latent_w = compute_latent_shape(height, width)
+        if generator is None:
+            latents = torch.randn((batch, channels, latent_h, latent_w), device=device, dtype=torch.float32)
+        else:
+            latents = randn_tensor((batch, channels, latent_h, latent_w), generator=generator, device=device, dtype=torch.float32)
+        latent = latents.cpu().numpy()
+        init_image_show = None
+        timesteps = TXT2IMG_TIMESTEPS
+        self_timesteps = None
+        step_index = None
+
+    latent = denoise_loop(
+        latent=latent,
+        prompt_embeds=prompt_embeds_npy,
+        time_inputs=time_input,
+        timesteps=timesteps,
+        unet_session=unet_session,
+        alphas_cumprod=alphas_cumprod,
+        final_alphas_cumprod=final_alphas_cumprod,
+        generator=generator,
+        noise_dtype=noise_dtype,
+        self_timesteps=self_timesteps,
+        step_index=step_index,
+    )
+
+    latent = latent / 0.18215
+    vae_decoder_inp_name = vae_decoder_session.get_inputs()[0].name
+    image = vae_decoder_session.run(None, {vae_decoder_inp_name: latent.astype(np.float32)})[0]
+
+    image = np.transpose(image, (0, 2, 3, 1)).squeeze(axis=0)
+    image_denorm = np.clip(image / 2 + 0.5, 0, 1)
+    image_uint8 = (image_denorm * 255).round().astype("uint8")
+    pil_image = Image.fromarray(image_uint8[:, :, :3])
+
+    grid_img = None
+    if is_img2img:
+        grid_img = make_image_grid([init_image_show, pil_image], rows=1, cols=2)
+
+    if save_path:
+        ensure_parent(save_path)
+        pil_image.save(save_path)
+        if grid_img is not None:
+            grid_path = os.path.splitext(save_path)[0] + "_grid.png"
+            ensure_parent(grid_path)
+            grid_img.save(grid_path)
+
+    return pil_image, grid_img, seed_used
+
+
+def gradio_generate(prompt: str,
+                    init_image: Optional[Image.Image],
+                    backend: str,
+                    isize: str,
+                    seed: Optional[float],
+                    model_dir: str):
+    try:
+        image, grid_img, seed_used = run_pipeline(
+            prompt=prompt,
+            model_dir=model_dir,
+            backend=backend,
+            isize=isize,
+            seed=int(seed) if seed not in (None, "") else None,
+            init_image=init_image,
+        )
+        return image, grid_img, f"{seed_used}"
+    except Exception as exc:  # pragma: no cover
+        warnings.warn(f"生成失败: {exc}")
+        return None, None, "生成失败"
+
+
+def launch_gradio(
+    default_model_dir: str = "./models",
+    default_backend: str = "axe",
+    default_isize: str = "512",
+    server_name: Optional[str] = None,
+    server_port: Optional[int] = None,
+    share: bool = False,
+):
+    # 先加载模型，若失败直接抛出，避免用户打开页面后才发现错误
+    _check_gradio_version()
+    print("[INIT] 正在预加载模型与时间输入...")
+    _preload_models(default_model_dir, default_backend, default_isize)
+    print("[INIT] 模型预加载完成")
+    title_text = "Stable Diffusion LCM Demo"
+    subtitle_text = f"分辨率 {default_isize}"
+    with gr.Blocks(title=title_text) as demo:
+        gr.Markdown(f"### {title_text}")
+        gr.Markdown(f"**{subtitle_text}**")
+        gr.HTML(
+            """
+            <style>
+            .fixed-img-container {height: 320px; display:flex; align-items:center; justify-content:center; overflow:hidden;}
+            .fixed-img-container img {max-height: 100%; max-width: 100%; object-fit: contain;}
+            .gradio-fullscreen img {max-height: none !important; width: auto !important; height: auto !important; object-fit: contain;}
+            </style>
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    lines=4,
+                    value="Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
+                    placeholder="输入提示词"
+                )
+                init_image = gr.Image(
+                    label="Init Image (可选)",
+                    type="pil",
+                    image_mode="RGB",
+                    elem_classes=["fixed-img-container"],
+                    show_fullscreen_button=True,
+                )
+                seed = gr.Number(label="随机种子 (-1 表示随机)", value=-1, precision=0)
+                seed_info = gr.Textbox(label="实际种子", value="-", interactive=False)
+                run_btn = gr.Button("生成", variant="primary")
+            with gr.Column(scale=1):
+                output_image = gr.Image(
+                    label="输出图像",
+                    elem_classes=["fixed-img-container"],
+                    show_fullscreen_button=True,
+                )
+                grid_image = gr.Image(
+                    label="对比图 (img2img)",
+                    elem_classes=["fixed-img-container"],
+                    show_fullscreen_button=True,
+                )
+
+        run_btn.click(
+            fn=gradio_generate,
+            inputs=[prompt, init_image, gr.State(default_backend), gr.State(default_isize), seed, gr.State(default_model_dir)],
+            outputs=[output_image, grid_image, seed_info],
+        )
+
+    app = demo.queue(max_size=4)
+
+    target_port = server_port or 7860
+    host_candidates = []
+    if server_name:
+        host_candidates.append(server_name)
+    host_candidates.extend(_list_host_ips())
+    printed = set()
+    print("可访问地址 (请任选其一):")
+    for ip in host_candidates:
+        if ip and ip not in printed:
+            printed.add(ip)
+            print(f"  http://{ip}:{target_port}")
+
+    app.launch(server_name=server_name, server_port=server_port, share=share)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for Stable Diffusion LCM")
+    parser.add_argument("--model_dir", type=str, default="./models", help="模型目录路径")
+    parser.add_argument("--backend", choices=["axe", "onnx"], default="axe", help="推理后端")
+    parser.add_argument("--isize", type=str, default="512x512", help="输出分辨率，单值或HxW，需为8的倍数")
+    parser.add_argument("--server_name", type=str, default="0.0.0.0", help="Gradio server_name，例如0.0.0.0")
+    parser.add_argument("--server_port", type=int, default=7860, help="Gradio server_port，例如7860")
+    parser.add_argument("--share", action="store_true", help="开启 Gradio share 链接")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    """
+    pip3 install gradio==5.42.0
+    python3 gradio_demo.py --model_dir models_1024x768 --isize 1024x768
+    """
+    args = get_args()
+    launch_gradio(
+        default_model_dir=args.model_dir,
+        default_backend=args.backend,
+        default_isize=args.isize,
+        server_name=args.server_name,
+        server_port=args.server_port,
+        share=args.share,
+    )