cli.py

#! /usr/bin/env python3
import torch
from diffusers import ZImagePipeline
from datetime import datetime
import os
import time
import warnings

# 禁用 transformers/diffusers 的进度条输出，避免干扰 Rich 进度条
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

# 禁用 transformers 的进度条
try:
    from transformers.utils import logging as transformers_logging
    transformers_logging.set_verbosity_error()
    transformers_logging.disable_progress_bar()
except ImportError:
    pass

# 禁用 warnings（可选）
warnings.filterwarnings("ignore")

# 美化输出和增强输入
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn
from rich.syntax import Syntax
from rich.text import Text
from prompt_toolkit import prompt
from prompt_toolkit.history import FileHistory
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.styles import Style

# 初始化
console = Console()
OUTPUT_DIR = "./outputs"
HISTORY_FILE = "./.zimage_history"  # 当前目录存储历史记录
MODEL_PATH = "./models/Z-Image-Turbo"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# 命令自动补全
command_completer = WordCompleter([
    '/set', '/info', '/quit', '/q', '/help', '/clear',
    'steps=', 'seed=',
    '512x512', '768x768', '1024x1024', '1280x1280',
], ignore_case=True)

# 输入框样式
prompt_style = Style.from_dict({
    'prompt': 'cyan bold',
})

def load_pipeline(model_path=MODEL_PATH):
    """加载并优化模型管道"""
    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]正在加载模型..."),
        BarColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=True,  # 结束后清理进度行，避免遗留控制字符如 [A
    ) as progress:
        task = progress.add_task("loading", total=None)
        start_time = time.time()
        torch.cuda.empty_cache()
        
        pipe = ZImagePipeline.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,  # 必须保留！虽然已弃用，但移除会导致 OOM
            # 注意：新版本提示使用 dtype，但 ZImagePipeline 可能还不支持，必须使用 torch_dtype
        )
        pipe.enable_model_cpu_offload()
        elapsed = time.time() - start_time
    
    console.print(f"[bold green]✓ 模型加载完成！[/] [dim]耗时: {elapsed:.2f}秒[/]\n")
    return pipe

def generate_image(pipe, prompt_text, width=1024, height=1024, steps=8, seed=None):
    """生成图像"""
    if seed is None:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    # 显示参数
    param_table = Table(show_header=False, box=None, padding=(0, 2))
    param_table.add_column(style="cyan")
    param_table.add_column(style="white")
    param_table.add_row("尺寸", f"{width}x{height}")
    param_table.add_row("步数", str(steps))
    param_table.add_row("种子", str(seed))
    
    console.print(Panel(param_table, title="[bold]生成参数[/]", border_style="blue"))
    
    # 显示提示词
    display_prompt = prompt_text[:100] + "..." if len(prompt_text) > 100 else prompt_text
    console.print(f"[bold]提示词:[/] [italic]{display_prompt}[/]\n")
    
    # 生成图像（带进度条）
    with Progress(
        SpinnerColumn(),
        TextColumn("[bold magenta]正在生成图像..."),
        BarColumn(complete_style="magenta"),
        TimeElapsedColumn(),
        console=console,
        transient=True,  # 结束后清理进度行，避免遗留控制字符
    ) as progress:
        task = progress.add_task("generating", total=None)
        start_time = time.time()
        
        image = pipe(
            prompt=prompt_text,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=0.0,
            generator=torch.Generator("cuda").manual_seed(seed)
        ).images[0]
        
        gen_time = time.time() - start_time
    
    # 保存图像
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{OUTPUT_DIR}/img_{timestamp}_{seed}.png"
    image.save(filename)
    
    # 显示结果
    console.print(f"\n[bold green]✓ 图像已保存:[/] [underline]{filename}[/]")
    console.print(f"[bold yellow]⏱ 生成耗时:[/] {gen_time:.2f}秒\n")
    
    return filename

def parse_settings(settings_str, current_width, current_height, current_steps, current_seed):
    """解析设置字符串"""
    width, height, steps, seed = current_width, current_height, current_steps, current_seed
    
    parts = settings_str.split()
    for part in parts:
        part_lower = part.lower()
        if 'x' in part_lower:
            try:
                w, h = part_lower.split('x')
                width, height = int(w), int(h)
            except:
                pass
        elif part_lower.startswith('steps='):
            try:
                steps = int(part.split('=')[1])
            except:
                pass
        elif part_lower.startswith('seed='):
            val = part.split('=')[1]
            seed = None if val.lower() in ['random', 'none', '随机'] else int(val)
    
    return width, height, steps, seed

def show_help():
    """显示帮助信息"""
    help_table = Table(title="命令帮助", show_header=True, header_style="bold cyan")
    help_table.add_column("命令", style="green")
    help_table.add_column("说明", style="white")
    help_table.add_column("示例", style="dim")
    
    help_table.add_row("提示词", "直接输入文字生成图像", "一只可爱的猫咪")
    help_table.add_row("/set", "修改生成参数", "/set 768x768 steps=4")
    help_table.add_row("/info", "查看当前参数", "/info")
    help_table.add_row("/clear", "清空历史记录", "/clear")
    help_table.add_row("/help", "显示帮助", "/help")
    help_table.add_row("/quit", "退出程序", "/quit 或 /q")
    
    console.print(help_table)
    console.print("\n[dim]提示: 按 ↑/↓ 键浏览输入历史，支持自动补全[/]\n")

def show_info(width, height, steps, seed):
    """显示当前参数"""
    info_table = Table(show_header=False, box=None)
    info_table.add_column(style="cyan bold")
    info_table.add_column(style="white")
    
    info_table.add_row("图像尺寸", f"{width} x {height}")
    info_table.add_row("推理步数", str(steps))
    info_table.add_row("随机种子", str(seed) if seed else "随机")
    info_table.add_row("输出目录", OUTPUT_DIR)
    
    console.print(Panel(info_table, title="[bold]当前参数[/]", border_style="cyan"))

def main():
    # 显示欢迎界面
    console.print(Panel.fit(
        "[bold magenta]Z-Image 图像生成器[/]\n"
        "[dim]交互模式 · 输入 /help 查看帮助[/]",
        border_style="magenta"
    ))
    console.print()
    
    # 加载模型
    pipe = load_pipeline()
    
    # 初始化输入历史
    history = FileHistory(HISTORY_FILE)
    
    # 默认参数
    width, height, steps, seed = 1024, 1024, 8, None
    
    while True:
        try:
            # 使用 prompt_toolkit 获取输入
            user_input = prompt(
                [('class:prompt', '🎨 提示词: ')],
                history=history,
                auto_suggest=AutoSuggestFromHistory(),
                completer=command_completer,
                style=prompt_style,
            ).strip()
            
            if not user_input:
                continue
            
            cmd = user_input.lower()
            
            # 处理命令
            if cmd in ['/quit', '/q', 'quit', 'exit']:
                console.print("[bold]再见！👋[/]")
                break
            
            if cmd == '/help':
                show_help()
                continue
            
            if cmd == '/info':
                show_info(width, height, steps, seed)
                continue
            
            if cmd == '/clear':
                os.remove(HISTORY_FILE) if os.path.exists(HISTORY_FILE) else None
                console.print("[green]✓ 历史记录已清空[/]")
                continue
            
            if cmd.startswith('/set '):
                settings = user_input[5:]
                width, height, steps, seed = parse_settings(settings, width, height, steps, seed)
                show_info(width, height, steps, seed)
                continue
            
            if cmd.startswith('/'):
                console.print(f"[red]未知命令: {user_input}[/]  输入 /help 查看帮助")
                continue
            
            # 生成图像
            generate_image(pipe, user_input, width, height, steps, seed)
            
        except KeyboardInterrupt:
            console.print("\n[bold]已中断，再见！👋[/]")
            break
        except EOFError:
            break
        except Exception as e:
            console.print(f"[bold red]错误:[/] {e}")
            continue

if __name__ == "__main__":
    main()