Diffusers.py

import time
import torch
import torch.nn.functional as F
import os
import argparse
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import piexif
from datetime import datetime
import gc
import sys
import tomesd
import threading
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import json
import traceback
import logging
import re
import random
import queue

# =============== 参数配置区 (方便后续修改) ===============
MODEL_NAME = "oneObsession_v16Noobai"
MODEL_DIR = "../ComfyUI/models/checkpoints"
OUTPUT_DIR = "../ComfyUI/output"
VPRED = False
WIDTH = 832
HEIGHT = 1216
CFG = 5
BASE_STEPS = 30
HIRES_UPSCALE_FACTOR = 1.5
HIRES_DENOISING_STRENGTH = 0.4
HIRES_STEPS = 20
PROMPT_ADD = "masterpiece,best quality,amazing quality,high quality,very awa,very aesthetic,newest,absurdres,highres,"
NEGATIVE_PROMPT = "worst quality,low quality,bad quality,worst aesthetic,ugly,old,early,lowres,worst detail,low details,jpeg artifacts,bad anatomy,bad hands,bad fingers,bad feet,"
BATCH_SIZE = 6
OUTPUT_FORMAT = "jpg"
# ======================================================

#os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
#os.environ["HF_HUB_OFFLINE"] = "1"

if torch.cuda.get_device_name().endswith("[ZLUDA]"):
    os.environ["DISABLE_ADDMM_CUDA_LT"] = "1"
    os.environ["ZLUDA_COMGR_LOG_LEVEL"] = "1"
    os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE"
    os.environ["FLASH_ATTENTION_TRITON_AMD_AUTOTUNE"] = "TRUE"
    os.environ["MIOPEN_FIND_MODE"] = "2"
    os.environ["MIOPEN_LOG_LEVEL"] = "3"
    torch.backends.cudnn.enabled = True
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(False)
    torch.backends.cuda.enable_mem_efficient_sdp = lambda _: None
    
    _topk = torch.topk
    def safe_topk(input: torch.Tensor, *args, **kwargs):
        device = input.device
        values, indices = _topk(input.cpu(), *args, **kwargs)
        return torch.return_types.topk((values.to(device), indices.to(device),))
    torch.topk = safe_topk

    try:
        import triton
        import triton.language as tl
        @triton.jit
        def _zluda_kernel_test(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
            pid = tl.program_id(axis=0)
            offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
            mask = offsets < n_elements
            x = tl.load(x_ptr + offsets, mask=mask)
            tl.store(y_ptr + offsets, x + 1, mask=mask)
        def _verify_triton() -> bool:
            try:
                x = torch.ones(64, device='cuda')
                y = torch.empty_like(x)
                _zluda_kernel_test[(1,)](x, y, x.numel(), BLOCK_SIZE=64)
                if torch.allclose(y, x + 1):
                    return True
                print("Triton kernel test failed (incorrect output)")
                return False
            except Exception as e:
                print(f"Triton test failed: {str(e)}")
                return False
        if _verify_triton():
            MEM_BUS_WIDTH = {
                "AMD Radeon RX 9070 XT": 256,
                "AMD Radeon RX 9070": 256,
                "AMD Radeon RX 9060 XT": 192,
                "AMD Radeon RX 7900 XTX": 384,
                "AMD Radeon RX 7900 XT": 320,
                "AMD Radeon RX 7900 GRE": 256,
                "AMD Radeon RX 7800 XT": 256,
                "AMD Radeon RX 7700 XT": 192,
                "AMD Radeon RX 7700": 192,
                "AMD Radeon RX 7650 GRE": 128,
                "AMD Radeon RX 7600 XT": 128,
                "AMD Radeon RX 7600": 128,
                "AMD Radeon RX 7500 XT": 96,
                "AMD Radeon RX 6950 XT": 256,
                "AMD Radeon RX 6900 XT": 256,
                "AMD Radeon RX 6800 XT": 256,
                "AMD Radeon RX 6800": 256,
                "AMD Radeon RX 6750 XT": 192,
                "AMD Radeon RX 6700 XT": 192,
                "AMD Radeon RX 6700": 160,
                "AMD Radeon RX 6650 XT": 128,
                "AMD Radeon RX 6600 XT": 128,
                "AMD Radeon RX 6600": 128,
                "AMD Radeon RX 6500 XT": 64,
                "AMD Radeon RX 6400": 64,
                "AMD Radeon 780M Graphics": 128,
            }
            _get_props = triton.runtime.driver.active.utils.get_device_properties
            def patched_props(device):
                props = _get_props(device)
                name = torch.cuda.get_device_name()[:-8]
                props["mem_bus_width"] = MEM_BUS_WIDTH.get(name, 128)
                return props
            triton.runtime.driver.active.utils.get_device_properties = patched_props
            try:
                sys.path.append(os.path.dirname(os.path.abspath(__file__)))
                from flash_attn_triton_amd import interface_fa
                original_sdpa = torch.nn.functional.scaled_dot_product_attention
                def amd_flash_wrapper(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
                    try:
                        if (query.shape[-1] <= 128 and 
                            attn_mask is None and 
                            query.dtype != torch.float32):
                            if scale is None:
                                scale = query.shape[-1] ** -0.5
                            return interface_fa.fwd(
                                query.transpose(1, 2),
                                key.transpose(1, 2),
                                value.transpose(1, 2),
                                None, None, dropout_p, scale,
                                is_causal, -1, -1, 0.0, False, None
                            )[0].transpose(1, 2)
                    except Exception as e:
                        print(f'Flash attention error: {str(e)}')
                    return original_sdpa(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
                torch.nn.functional.scaled_dot_product_attention = amd_flash_wrapper
                print("Flash attention enabled successfully")
            except ImportError:
                print("Flash attention components not installed")
            except Exception as e:
                print(f"Flash attention setup failed: {str(e)}")
        else:
            print("Triton available but failed verification")
    except ImportError:
        print("Triton not installed")
    except Exception as e:
        print(f"Triton initialization failed: {str(e)}")

app = Flask(__name__)
CORS(app, resources={r"/api/*": {"origins": "*"}})
logging.getLogger('werkzeug').setLevel(logging.ERROR)

GENERATION_STATUS = {
    "is_running": False,
    "progress": 0,
    "current_step": "",
    "total_images": 0,
    "processed_images": 0,
    "error": None,
    "last_error_traceback": None,
    "queue_size": 0
}

TASK_QUEUE = queue.Queue()
QUEUE_LOCK = threading.Lock()
CURRENT_ACTIVE_CONFIG = None

def queue_processor():
    global CURRENT_ACTIVE_CONFIG
    while True:
        try:
            task = TASK_QUEUE.get()
            config_dict = task['config']
            prompt_lines = task['prompts']

            with QUEUE_LOCK:
                current_config_snapshot = CURRENT_ACTIVE_CONFIG.copy() if CURRENT_ACTIVE_CONFIG else {}

            merged_config = {**current_config_snapshot, **config_dict}
            with QUEUE_LOCK:
                CURRENT_ACTIVE_CONFIG = config_dict.copy()

            class Config:
                pass
            config = Config()
            for key, value in merged_config.items():
                snake_case_key = re.sub(r'(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])', '_', key).lower()
                setattr(config, snake_case_key, value)

            if not hasattr(config, 'model_dir'):
                config.model_dir = MODEL_DIR
            if not hasattr(config, 'output_dir'):
                config.output_dir = OUTPUT_DIR
            if not hasattr(config, 'prompt_add'):
                config.prompt_add = PROMPT_ADD
            if not hasattr(config, 'negative'):
                config.negative = NEGATIVE_PROMPT
            if not hasattr(config, 'hires_batch'):
                config.hires_batch = max(1, config.batch // 2)
            config.width = (config.width + 32) // 64 * 64
            config.height = (config.height + 32) // 64 * 64
            if hasattr(config, 'seed') and (config.seed == "" or config.seed is None):
                config.seed = None
            elif hasattr(config, 'seed') and config.seed != None:
                try:
                    config.seed = int(config.seed)
                except:
                    config.seed = None

            config.prompt = "\\n".join(prompt_lines)

            GENERATION_STATUS["is_running"] = True
            GENERATION_STATUS["progress"] = 5
            GENERATION_STATUS["current_step"] = "初始化..."
            GENERATION_STATUS["total_images"] = len(prompt_lines)
            GENERATION_STATUS["processed_images"] = 0
            GENERATION_STATUS["error"] = None
            GENERATION_STATUS["last_error_traceback"] = None

            run_generation(config)

            TASK_QUEUE.task_done()

            if not TASK_QUEUE.empty():
                GENERATION_STATUS["current_step"] = "队列中还有任务，准备下一批..."
                time.sleep(2)
            else:
                GENERATION_STATUS["is_running"] = False
                GENERATION_STATUS["current_step"] = "空闲 - 等待新任务"
                GENERATION_STATUS["progress"] = 0
                GENERATION_STATUS["queue_size"] = 0

        except Exception as e:
            error_msg = str(e)
            traceback_str = traceback.format_exc()
            print(f"❌ 队列处理器出错: {error_msg}")
            print(traceback_str)
            GENERATION_STATUS["error"] = error_msg
            GENERATION_STATUS["last_error_traceback"] = traceback_str
            GENERATION_STATUS["is_running"] = False
            if not TASK_QUEUE.empty():
                TASK_QUEUE.task_done()

def cleanup_resources():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    gc.collect()

def pipeline_config(pipeline, config):
    prediction_type = "v_prediction" if config.vpred else "epsilon"
    rescale_betas = config.vpred
    
    pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
        pipeline.scheduler.config,
        prediction_type=prediction_type,
        rescale_betas_zero_snr=rescale_betas
    )
    
    pipeline.enable_vae_slicing()
    pipeline.enable_vae_tiling()
    
    return pipeline

def save_image_with_metadata(img, filename, prompt, seed_value, config):
    metadata_str = (
        f"Prompt: {prompt}\n"
        f"Negative prompt: {config.negative}\n"
        f"Seed: {seed_value}\n"
        f"Model: {config.model_name}\n"
        f"CFG scale: {config.cfg}\n"
        f"Steps: Base: {config.steps}, HiRes: {config.hires_steps if config.upscale > 1.0 else 'disabled'}\n"
        f"HiRes upscale: {config.upscale if config.upscale > 1.0 else 'disabled'} (denoising: {config.denoise if config.upscale > 1.0 else 'disabled'})"
    )
    
    if filename.lower().endswith('.png'):
        pnginfo = PngInfo()
        for line in metadata_str.split('\n'):
            key, value = line.split(':', 1)
            pnginfo.add_text(key.strip(), value.strip())
        img.save(filename, pnginfo=pnginfo)
    elif filename.lower().endswith('.jpg') or filename.lower().endswith('.jpeg'):
        exif_dict = {"Exif": {}}
        exif_dict["0th"] = {
            piexif.ImageIFD.Model: config.model_name.encode('utf-8'),
            piexif.ImageIFD.Software: b"StableDiffusionXL Generator",
        }
        exif_dict["Exif"][piexif.ExifIFD.UserComment] = metadata_str.encode('utf-8')
        exif_bytes = piexif.dump(exif_dict)
        img.save(filename, "JPEG", quality=95, exif=exif_bytes)
    else:
        img.save(filename)

def run_generation(config):
    start_time = time.time()
    
    if hasattr(config, 'prompt') and config.prompt is not None:
        prompts = [line.strip() + config.prompt_add for line in config.prompt.split('\n') if line.strip()]
        print(f"📌 使用命令行指定的 {len(prompts)} 条提示词")
    else:
        prompt_file = os.path.join(config.output_dir, "@prompt.txt")
        if not os.path.exists(prompt_file) or not os.path.getsize(prompt_file):
            print(f"❌ 错误: 无效的 prompt 文件: {prompt_file}")
            return
        
        with open(prompt_file, "r", encoding="utf-8") as f:
            prompts = [line.strip() + config.prompt_add for line in f if line.strip()]
        print(f"📌 从文件加载 {len(prompts)} 条提示词: {prompt_file}")
    
    model_path = os.path.join(config.model_dir, f"{config.model_name}.safetensors")
    if not os.path.exists(model_path):
        print(f"❌ 错误: 模型文件不存在: {model_path}")
        return
    
    hires_fix_enabled = config.upscale > 1.0
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    total_base = 0
    total_hires = 0
    success_count = 0
    if config.seed is None:
        seed_value = random.randint(0, 2**32-1)
    else:
        seed_value = config.seed

    print(f"🎨 [{config.model_name}]开始生成 {len(prompts)} 张图像 ({config.width}x{config.height})...")
    print(f"📦 基础批处理大小: {config.batch}")
    if hires_fix_enabled:
        print(f"🔄 HiRes批处理大小: {config.hires_batch}, 放大因子: {config.upscale}")
    else:
        print("❌ HiRes修复已禁用 (放大因子 <= 1.0)")
    
    base_pipeline = StableDiffusionXLPipeline.from_single_file(
        model_path,
        torch_dtype=torch.float16,
        use_safetensors=True
    )
    base_pipeline = pipeline_config(base_pipeline, config)
    tomesd.apply_patch(base_pipeline, ratio=0.5)
    base_pipeline.enable_model_cpu_offload()
    
    for i in range(0, len(prompts), config.batch):
        if GENERATION_STATUS["is_running"]:
            GENERATION_STATUS["processed_images"] = i
            GENERATION_STATUS["progress"] = 15 + (i / len(prompts)) * 85
            GENERATION_STATUS["current_step"] = f"生成基础图像 {i+1}/{len(prompts)}"
        
        batch = prompts[i:i+config.batch]
        batch_size = len(batch)
        
        base_start = time.time()
        
        if hires_fix_enabled:
            base_images = base_pipeline(
                prompt=batch,
                negative_prompt=[config.negative] * batch_size,
                width=config.width,
                height=config.height,
                guidance_scale=config.cfg,
                num_inference_steps=config.steps,
                num_images_per_prompt=1,
                generator=torch.Generator(device="cpu").manual_seed(seed_value),
                output_type="latent",
            ).images
            
            total_base += time.time() - base_start
            hires_start = time.time()
            
            hires_pipeline = StableDiffusionXLImg2ImgPipeline(
                vae=base_pipeline.vae,
                text_encoder=base_pipeline.text_encoder,
                text_encoder_2=base_pipeline.text_encoder_2,
                tokenizer=base_pipeline.tokenizer,
                tokenizer_2=base_pipeline.tokenizer_2,
                unet=base_pipeline.unet,
                scheduler=base_pipeline.scheduler,
                force_zeros_for_empty_prompt=True
            )
            hires_pipeline = pipeline_config(hires_pipeline, config)
            tomesd.apply_patch(hires_pipeline, ratio=0.5)
            
            latent_w = int(config.width / 8 * config.upscale)
            latent_h = int(config.height / 8 * config.upscale)
            scale_w = int(config.width * config.upscale)
            scale_h = int(config.height * config.upscale)
            
            for j in range(0, batch_size, config.hires_batch):
                hires_batch_end = min(j + config.hires_batch, batch_size)
                hires_batch_size_current = hires_batch_end - j
                
                if GENERATION_STATUS["is_running"]:
                    current_idx = i + j
                    GENERATION_STATUS["processed_images"] = current_idx
                    GENERATION_STATUS["progress"] = 15 + (current_idx / len(prompts)) * 85
                    GENERATION_STATUS["current_step"] = f"处理HiRes {current_idx+1}/{len(prompts)}"
                
                current_batch_latents = base_images[j:hires_batch_end]
                upscaled = F.interpolate(
                    current_batch_latents,
                    size=(latent_h, latent_w),
                    mode="bicubic",
                )
                current_batch_prompts = batch[j:hires_batch_end]
                
                hires_imgs = hires_pipeline(
                    prompt=current_batch_prompts,
                    negative_prompt=[config.negative] * hires_batch_size_current,
                    image=upscaled,
                    width=scale_w,
                    height=scale_h,
                    guidance_scale=config.cfg,
                    num_inference_steps=config.hires_steps,
                    strength=config.denoise,
                    generator=torch.Generator(device="cpu").manual_seed(seed_value)
                ).images
                
                for k, img in enumerate(hires_imgs):
                    idx = i + j + k + 1
                    filename = os.path.join(config.output_dir, f"{config.model_name}-{timestamp}-{idx:03d}.{config.output_format}")
                    save_image_with_metadata(img, filename, current_batch_prompts[k], seed_value, config)
                    print(f"  ✅ #{idx:03d} 生成完成")
                    success_count += 1
                    
                if 'upscaled' in locals(): del upscaled
                if 'hires_imgs' in locals(): del hires_imgs
                cleanup_resources()
                
            total_hires += time.time() - hires_start
            
            if 'base_images' in locals(): del base_images
            if 'hires_pipeline' in locals(): del hires_pipeline
            cleanup_resources()
            
        else:
            base_images = base_pipeline(
                prompt=batch,
                negative_prompt=[config.negative] * batch_size,
                width=config.width,
                height=config.height,
                guidance_scale=config.cfg,
                num_inference_steps=config.steps,
                num_images_per_prompt=1,
                generator=torch.Generator(device="cpu").manual_seed(seed_value)
            ).images
            
            for j, img in enumerate(base_images):
                idx = i + j + 1
                filename = os.path.join(config.output_dir, f"{config.model_name}-{timestamp}-{idx:03d}.{config.output_format}")
                save_image_with_metadata(img, filename, batch[j], seed_value, config)
                print(f"  ✅ #{idx:03d} 生成完成")
                success_count += 1
            total_base += time.time() - base_start
            if 'base_images' in locals(): del base_images
            cleanup_resources()
    
    if 'base_pipeline' in locals(): del base_pipeline
    cleanup_resources()

    total_time = time.time() - start_time
    if success_count > 0:
        print("\n" + "="*50)
        print(f"🎉 [{config.model_name}] 生成完成! 总耗时: {total_time:.1f} 秒")
        print(f"📊 成功生成: {success_count}/{len(prompts)} 张图像")
        print(f"⏱️  平均速度: {total_time/success_count:.1f} 秒/图")
        if hires_fix_enabled:
            print(f"   - 基础图: {total_base/success_count:.1f} 秒/图")
            print(f"   - HiRes: {total_hires/success_count:.1f} 秒/图")
        print("="*50)
    
    if GENERATION_STATUS["is_running"]:
        GENERATION_STATUS["progress"] = 100
        GENERATION_STATUS["current_step"] = "生成完成!"
        GENERATION_STATUS["processed_images"] = len(prompts)
        time.sleep(1)

@app.route('/api/generate', methods=['POST'])
def api_generate():
    try:
        data = request.json
        if not data.get('prompt'):
            return jsonify({"error": "必须提供提示词"}), 400

        prompt_text = data.get('prompt', '')
        prompt_lines = [line.strip() for line in prompt_text.split('\\n') if line.strip()]

        if not prompt_lines:
            return jsonify({"error": "提示词列表为空"}), 400

        batch_size = data.get('batchSize', BATCH_SIZE)
        for i in range(0, len(prompt_lines), batch_size):
            batch_prompts = prompt_lines[i:i + batch_size]

            task = {
                'config': data.copy(),
                'prompts': batch_prompts
            }

            with QUEUE_LOCK:
                global CURRENT_ACTIVE_CONFIG
                CURRENT_ACTIVE_CONFIG = data.copy()

            TASK_QUEUE.put(task)
        
        GENERATION_STATUS["queue_size"] = TASK_QUEUE.qsize()

        return jsonify({
            "status": "任务已加入队列",
            "queued_batches": (len(prompt_lines) + batch_size - 1) // batch_size,
            "total_prompts": len(prompt_lines),
            "queue_size": GENERATION_STATUS["queue_size"]
        })

    except Exception as e:
        return jsonify({
            "error": str(e),
            "traceback": traceback.format_exc()
        }), 500

@app.route('/api/status', methods=['GET'])
def api_status():
    return jsonify({
        "isRunning": GENERATION_STATUS["is_running"],
        "progress": GENERATION_STATUS["progress"],
        "currentStep": GENERATION_STATUS["current_step"],
        "totalImages": GENERATION_STATUS["total_images"],
        "processedImages": GENERATION_STATUS["processed_images"],
        "error": GENERATION_STATUS["error"],
        "lastErrorTraceback": GENERATION_STATUS["last_error_traceback"],
        "queueSize": GENERATION_STATUS["queue_size"]
    })

def run_flask_server():
    app.run(host='0.0.0.0', port=8188, threaded=True, use_reloader=False)

if __name__ == "__main__":
    if "--api" in sys.argv:
        print("启动API服务器...")
        flask_thread = threading.Thread(target=run_flask_server)
        flask_thread.daemon = True
        flask_thread.start()
        print("API服务器已启动")
        
        queue_thread = threading.Thread(target=queue_processor)
        queue_thread.daemon = True
        queue_thread.start()
        print("队列处理器已启动")
        
        print("按Ctrl+C停止服务器")
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("\n服务器已停止")
    else:
        parser = argparse.ArgumentParser(description='Stable Diffusion XL 图像生成工具')
        parser.add_argument('--model-name', type=str, default=MODEL_NAME,
                            help='模型名称 (默认: %(default)s)')
        parser.add_argument('--model-dir', type=str, default=MODEL_DIR,
                            help='模型目录 (默认: %(default)s)')
        parser.add_argument('--output-dir', type=str, default=OUTPUT_DIR,
                            help='输出目录 (默认: %(default)s)')
        parser.add_argument('--vpred', action='store_true', default=VPRED,
                            help='启用v_prediction模式 (默认: epsilon模式)')
        parser.add_argument('--seed', type=int, default=None,
                            help='固定种子值 (默认: 随机)')
        parser.add_argument('--width', type=int, default=WIDTH,
                            help='图像宽度 (默认: %(default)s)')
        parser.add_argument('--height', type=int, default=HEIGHT,
                            help='图像高度 (默认: %(default)s)')
        parser.add_argument('--cfg', type=float, default=CFG,
                            help='分类器自由引导尺度 (默认: %(default)s)')
        parser.add_argument('--steps', type=int, default=BASE_STEPS,
                            help='基础生成步数 (默认: %(default)s)')
        parser.add_argument('--upscale', type=float, default=HIRES_UPSCALE_FACTOR,
                            help='HiRes 放大因子 (默认: %(default)s，<=1时禁用HiRes)')
        parser.add_argument('--denoise', type=float, default=HIRES_DENOISING_STRENGTH,
                            help='HiRes 去噪强度 (默认: %(default)s)')
        parser.add_argument('--hires-steps', type=int, default=HIRES_STEPS,
                            help='HiRes 生成步数 (默认: %(default)s)')
        parser.add_argument('--prompt-add', type=str, default=PROMPT_ADD,
                            help='附加的正向提示词 (默认: %(default)s)')
        parser.add_argument('--negative', type=str, default=NEGATIVE_PROMPT,
                            help='负面提示词 (默认: %(default)s)')
        parser.add_argument('--prompt', type=str, default=None,
                            help='直接指定prompt，使用\\n分割同一批次的不同prompt')
        parser.add_argument('--batch', type=int, default=BATCH_SIZE,
                            help='基础批处理大小 (默认: %(default)s)')
        parser.add_argument('--output-format', type=str, default=OUTPUT_FORMAT, choices=['png', 'jpg', 'jpeg'],
                            help='输出文件格式 (默认: %(default)s)')
        args = parser.parse_args()
        args.width = (args.width + 32) // 64 * 64
        args.height = (args.height + 32) // 64 * 64
        args.hires_batch = max(1, args.batch // 2)
        
        run_generation(args)