Stable Diffusion 离线下载、配置修复与作业自动追踪

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import sys
import time

def main():
    parser = argparse.ArgumentParser(description="超算纯净全自动离线版 SD 推理后端 (优化版)")
    parser.add_argument('--prompt', type=str, required=True, help="正向提示词")
    parser.add_argument('--negative_prompt', type=str, default='', help="反向提示词")
    parser.add_argument('--output_dir', type=str, default='./output_images', help="输出目录")
    parser.add_argument('--model_path', type=str, required=True, help="本地模型路径")
    parser.add_argument('--steps', type=int, default=28, help="渲染步数")
    parser.add_argument('--cfg', type=float, default=7.5, help="CFG 引导系数")
    parser.add_argument('--seed', type=int, default=-1, help="随机种子 (-1 表示随机产生)")
    parser.add_argument('--width', type=int, default=512, help="图片宽度 (默认 512)")
    parser.add_argument('--height', type=int, default=512, help="图片高度 (默认 512)")
    parser.add_argument('--num_images', type=int, default=1, help="批量生成图片张数 (默认 1)")
    parser.add_argument('--scheduler', type=str, default='euler_a', 
                        choices=['euler_a', 'euler', 'dpm++_2m', 'dpm++_2m_sde', 'ddim', 'pndm'],
                        help="采样器/调度器选择 (默认 euler_a)")
    parser.add_argument('--vae_fp32', action='store_true', default=False, 
                        help="强制 VAE 模块使用 FP32 精度以规避黑图/灰白图问题")
    parser.add_argument('--low_vram', action='store_true', default=False,
                        help="启用低显存优化 (Attention Slicing)，会牺牲 20%~40% 的速度，显存充足时切勿开启")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # 延迟加载重度依赖库，使 --help 及参数解析可以秒开
    print(">> [DEBUG] [1/6] 正在初始化 Python 环境，加载 PyTorch 核心库 (共享存储网络读取较慢，请稍候)...", flush=True)
    t_start = time.perf_counter()
    import torch
    print(f">> [DEBUG] [2/6] PyTorch 载入成功！当前 CUDA (DCU) 状态: {torch.cuda.is_available()} | 耗时: {time.perf_counter() - t_start:.2f}s", flush=True)

    print(">> [DEBUG] [3/6] 正在载入 diffusers 核心组件...", flush=True)
    t_diff = time.perf_counter()
    from diffusers import StableDiffusionPipeline
    print(f">> [DEBUG] [4/6] 核心组件载入完成！准备解析参数... | 耗时: {time.perf_counter() - t_diff:.2f}s", flush=True)

    # 锁定运行设备（由超算系统隔离分配）
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[INFO] 当前运行设备锁定为: {device}", flush=True)

    print(f"[INFO] [5/6] 正在从存储节点加载本地模型权重 (大文件读取，请耐心等待)...", flush=True)
    t_load = time.perf_counter()
    
    # 显式指定低版本兼容组件并完全关闭安全检查器（防止额外加载安全检测器权重占用显存）
    from transformers import CLIPFeatureExtractor
    feature_extractor = CLIPFeatureExtractor.from_pretrained(f"{args.model_path}/feature_extractor")
    
    pipe = StableDiffusionPipeline.from_pretrained(
        args.model_path, 
        # 从下方链接下载与 diffusers v0.21.4 匹配的 lpw 脚本放置于同级目录下：
        # https://github.com/huggingface/diffusers/blob/v0.21.4/examples/community/lpw_stable_diffusion.py
        custom_pipeline="./lpw_stable_diffusion.py",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_safetensors=True,
        feature_extractor=feature_extractor,
        safety_checker=None,                # 彻底禁用安全检查器模块，减少显存消耗和载入时间
        requires_safety_checker=False        
    )
    
    # 根据参数灵活配置 Scheduler
    scheduler_name = args.scheduler.lower()
    if scheduler_name == "euler_a":
        from diffusers import EulerAncestralDiscreteScheduler
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    elif scheduler_name == "euler":
        from diffusers import EulerDiscreteScheduler
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
    elif scheduler_name == "dpm++_2m":
        from diffusers import DPMSolverMultistepScheduler
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    elif scheduler_name == "dpm++_2m_sde":
        from diffusers import DPMSolverMultistepScheduler
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
    elif scheduler_name == "ddim":
        from diffusers import DDIMScheduler
        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
    elif scheduler_name == "pndm":
        from diffusers import PNDMScheduler
        pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)

    print(f"[INFO] 权重文件读取并解析完成。| 耗时: {time.perf_counter() - t_load:.2f}s", flush=True)
    
    print(f"[INFO] [6/6] 开始将模型分块搬运至显卡内存...", flush=True)
    t_to_device = time.perf_counter()
    pipe = pipe.to(device)
    if device == "cuda":
        torch.cuda.synchronize()  # 强行等待搬运彻底结束
    print(f"[INFO] 模型上卡成功！| 耗时: {time.perf_counter() - t_to_device:.2f}s", flush=True)

    # 显存优化策略配置
    vae_fp32 = args.vae_fp32
    if device == "cuda":
        if args.low_vram:
            print("[INFO] ⚡ 已开启低显存优化模式 (Attention Slicing)，将牺牲部分运算速度以降低显存占用。", flush=True)
            pipe.enable_attention_slicing()
        else:
            print("[INFO] 🚀 显存充足，采用满血极速模式运行（未启用 Attention Slicing）。", flush=True)
        
        # 动漫模型精度自动对齐策略（或手动强制开启）
        if not vae_fp32:
            model_tag = args.model_path.lower()
            if "anything" in model_tag or "counterfeit" in model_tag:
                vae_fp32 = True
                print("[INFO] 🛡️ 检测到二次元动漫模型，已自动触发开启 VAE FP32 精度对齐，防止灰沙图！", flush=True)
        
        if vae_fp32:
            print("[INFO] 🛡️ VAE 模块已强制升级至 FP32 精度，保证解码色域正常。", flush=True)
            pipe.vae.to(dtype=torch.float32)

    # 随机种子处理
    if args.seed == -1:
        import random
        base_seed = random.randint(0, 2**32 - 1)
    else:
        base_seed = args.seed
    
    print("==========================================", flush=True)
    print(f"[INFO] 基础随机种子: {base_seed} | 计划出图数量: {args.num_images}", flush=True)
    print("==========================================", flush=True)

    print("[INFO] 海光核心开始爆轰，正在解构潜空间...", flush=True)
    print("[INFO] 💡 提示：海光 DCU 在执行第 1 步时会默默编译 MIOpen 算子，耗时较长，请观察下方实时步数计数：", flush=True)

    # 开始循环渲染多张图片，最大化复用已载入的模型，规避超算排队和 PyTorch/Model 加载时间瓶颈
    for idx in range(args.num_images):
        current_seed = (base_seed + idx) & 0xFFFFFFFF
        generator = torch.Generator(device=device).manual_seed(current_seed)
        
        print(f"\n[INFO] [图片 {idx + 1}/{args.num_images}] 开始渲染，随机种子 Seed: {current_seed}", flush=True)
        
        # 定义渲染进度回调函数
        def step_callback(step: int, timestep: int, latents: torch.FloatTensor):
            print(f" -> [渲染雷达] 正在执行渲染: 图片 [{idx + 1}/{args.num_images}] | 步数 [{step + 1}/{args.steps}] | 戳记: {time.strftime('%H:%M:%S', time.localtime())}", flush=True)
        
        with torch.inference_mode():
            t_infer = time.perf_counter()
            # 1. 强行拉起管线进行潜空间扩散
            output = pipe(
                prompt=args.prompt,
                negative_prompt=args.negative_prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.cfg,
                generator=generator,
                width=args.width,
                height=args.height,
                output_type="latent",
                callback=step_callback,
                callback_steps=1
            )
            latents = output.images if hasattr(output, "images") else output[0]
            
            # 🔥【第一道硬件同步令】：斩断扩散流积压延迟
            if device == "cuda":
                torch.cuda.synchronize()  
            print(f"[INFO] [图片 {idx + 1}/{args.num_images}] 潜空间扩散生成完毕。| 纯推理耗时: {time.perf_counter() - t_infer:.2f}s", flush=True)
            
            print("[INFO] 正在进入 VAE 图像解码与矩阵转换阶段...", flush=True)
            t_vae = time.perf_counter()
            
            # 如果是 FP32 解码，强制把输入 VAE 的 latents 数据也提升到 FP32 精度对齐
            if vae_fp32 and device == "cuda":
                latents = latents.to(dtype=torch.float32)
                
            # 2. 用全面升级后的 FP32 精度进行 VAE 解码
            image_tensor = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
            
            # 🔥【第二道硬件同步令】：锁定 VAE 显存矩阵，防止异步流抢跑
            if device == "cuda":
                torch.cuda.synchronize()
            print(f"  -> [VAE明细] VAE网络解码矩阵完成。| 耗时: {time.perf_counter() - t_vae:.2f}s", flush=True)
            
            # 3. 安全后处理与内存回收
            t_post = time.perf_counter()
            image_tensor = (image_tensor / 2 + 0.5).clamp(0, 1)
            
            # 🔥【第三道硬件同步令】：拦截 D2H 跨节点搬运，防止后台总线死锁
            if device == "cuda":
                torch.cuda.synchronize()
                
            image_tensor = image_tensor.cpu()
            image_tensor = image_tensor.permute(0, 2, 3, 1).float().numpy()[0]
            
            # 4. 转换回标准的 PIL 图片对象
            image = pipe.numpy_to_pil(image_tensor)[0]
            print(f"[INFO] 后处理与 PIL 对象转换全部搞定。| 耗时: {time.perf_counter() - t_post:.2f}s", flush=True)

        # 保存落地
        output_filename = f"sd_seed_{current_seed}.png"
        output_path = os.path.join(args.output_dir, output_filename)
        image.save(output_path)
        
        print("==========================================", flush=True)
        print(f"🎉 炼丹成功！图片已安全落地 -> {output_path}", flush=True)
        print("==========================================", flush=True)
        
        # 释放每轮循环的临时显存，防止显存积攒发生 OOM
        if device == "cuda":
            import gc
            gc.collect()
            torch.cuda.empty_cache()

if __name__ == '__main__':
    main()