app.py

#!/usr/bin/env python3
"""
Zero123++ 工程六视图生成器
通过旋转输入图片来模拟不同视角
"""

import gradio as gr
import torch
from PIL import Image
import numpy as np
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler

# 全局变量存储 pipeline
pipeline = None

# 工程六视图配置
# Zero123++ 固定输出 6 个视图,我们选择最接近正交视图的角度
# 方位角: 30, 90, 150, 210, 270, 330
# 我们使用: 90° (右), 210° (后偏左), 330° (前偏右)
ENGINEERING_VIEWS = {
    "主视图": {"rotate_input": 0, "select_index": 0, "position": (0, 0)},     # 使用原始输入
    "右视图": {"rotate_input": 0, "select_index": 1, "position": (1, 0)},     # 90度视角
    "后视图": {"rotate_input": 0, "select_index": 3, "position": (2, 0)},     # 210度视角
    "左视图": {"rotate_input": 0, "select_index": 4, "position": (0, 1)},     # 270度视角
    "俯视图": {"rotate_input": -90, "select_index": 0, "position": (1, 1)},   # 旋转输入-90度
    "底视图": {"rotate_input": 90, "select_index": 0, "position": (2, 1)},    # 旋转输入+90度
}

def load_model():
    """加载 Zero123++ 模型"""
    global pipeline

    if pipeline is not None:
        return

    print("正在加载 Zero123++ 模型...")

    # 检查 CUDA 可用性
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    try:
        # 加载 Zero123++ pipeline
        pipeline = DiffusionPipeline.from_pretrained(
            "sudo-ai/zero123plus-v1.2",
            custom_pipeline="sudo-ai/zero123plus-pipeline",
            torch_dtype=dtype
        )

        # 设置调度器
        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
            pipeline.scheduler.config,
            timestep_spacing='trailing'
        )

        pipeline.to(device)

        # 启用内存优化
        if torch.cuda.is_available():
            try:
                pipeline.enable_attention_slicing()
                pipeline.enable_vae_slicing()
            except:
                pass

        print(f"✓ 模型加载完成 (设备: {device})")
    except Exception as e:
        print(f"错误: 无法加载 Zero123++ - {e}")
        raise

def rotate_image(image, angle):
    """旋转图像"""
    if angle == 0:
        return image
    elif angle == 90:
        return image.rotate(-90, expand=True)
    elif angle == -90:
        return image.rotate(90, expand=True)
    elif angle == 180:
        return image.rotate(180, expand=True)
    return image

def generate_multiview(input_image):
    """
    生成多视图输出

    输入:
        input_image: PIL Image

    输出:
        PIL Image with 6 views (2x3 grid)
    """
    global pipeline

    if pipeline is None:
        load_model()

    # 预处理输入图像
    img = input_image.resize((320, 320), Image.LANCZOS)

    # 运行推理生成6个视图
    result = pipeline(
        img,
        num_inference_steps=75
    ).images[0]

    # Zero123++ 输出是 2x3 的网格,需要拆分
    # 输出尺寸应该是 320*3 x 320*2 = 960x640
    result_w, result_h = result.size
    view_w = result_w // 3
    view_h = result_h // 2

    # 提取6个视图
    views = []
    for row in range(2):
        for col in range(3):
            x = col * view_w
            y = row * view_h
            view = result.crop((x, y, x + view_w, y + view_h))
            views.append(view)

    return views

def process_image(input_image, progress=gr.Progress()):
    """
    处理输入图像,生成工程六视图

    输入:
        input_image: PIL Image
        progress: Gradio Progress 跟踪器

    输出:
        result_image: PIL Image (六视图合成图)
    """
    if input_image is None:
        return None

    try:
        # 确保模型已加载
        load_model()

        # 预处理 - 转为正方形
        img = input_image
        if img.size[0] != img.size[1]:
            size = min(img.size)
            img = img.crop((
                (img.size[0] - size) // 2,
                (img.size[1] - size) // 2,
                (img.size[0] + size) // 2,
                (img.size[1] + size) // 2
            ))

        progress(0.1, desc="生成水平视图...")
        # 生成水平方向的视图 (主/右/后/左)
        horizontal_views = generate_multiview(img)

        progress(0.5, desc="生成俯视图...")
        # 旋转输入-90度生成俯视图
        img_rotated_up = rotate_image(img, -90)
        top_views = generate_multiview(img_rotated_up)

        progress(0.8, desc="生成底视图...")
        # 旋转输入+90度生成底视图
        img_rotated_down = rotate_image(img, 90)
        bottom_views = generate_multiview(img_rotated_down)

        # 组合成工程六视图
        # 布局: 主/右/后
        #       左/俯/底
        view_size = 320
        combined = Image.new('RGB', (view_size * 3, view_size * 2))

        # 第一行: 主视图(0), 右视图(90), 后视图(210)
        combined.paste(horizontal_views[0], (0 * view_size, 0))          # 主视图 (30度)
        combined.paste(horizontal_views[1], (1 * view_size, 0))          # 右视图 (90度)
        combined.paste(horizontal_views[3], (2 * view_size, 0))          # 后视图 (210度)

        # 第二行: 左视图(270), 俯视图, 底视图
        combined.paste(horizontal_views[4], (0 * view_size, view_size))  # 左视图 (270度)
        combined.paste(top_views[0], (1 * view_size, view_size))         # 俯视图 (旋转输入-90)
        combined.paste(bottom_views[0], (2 * view_size, view_size))      # 底视图 (旋转输入+90)

        progress(1.0, desc="完成!")
        print("✓ 所有视图生成完成")
        return combined

    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"处理失败: {str(e)}")

# 创建 Gradio 界面
def create_demo():
    with gr.Blocks(title="Zero123++ 工程六视图生成器") as demo:
        gr.Markdown("""
        # Zero123++ 工程六视图生成器

        将单张主视图转换为工程六视图

        **输入要求:**
        - 建议使用正方形图片
        - 推荐分辨率 >= 320x320
        - 脚本会自动裁剪非正方形图片

        **输出说明:**
        生成工程六视图,排列为 2 行 3 列:

        | 主视图 | 右视图 | 后视图 |
        |-------|-------|-------|
        | 左视图 | 俯视图 | 底视图 |

        **技术原理:**
        - 使用 Zero123++ v1.2 模型
        - 通过 3 次推理 + 旋转输入实现工程六视图
        - 每次推理约 30-60 秒,总计约 2-3 分钟
        """)

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="输入主视图",
                    type="pil",
                    height=400
                )

                generate_btn = gr.Button("生成工程六视图", variant="primary", size="lg")

                gr.Markdown("""
                **注意:**
                - 需要运行 3 次推理(水平+俯视+底视)
                - 总耗时约 2-3 分钟
                - 请耐心等待
                """)

            with gr.Column():
                output_image = gr.Image(
                    label="工程六视图输出",
                    type="pil",
                    height=400
                )

        gr.Markdown("""
        ### 视角说明

        | 视图 | 方法 | 说明 |
        |-----|------|------|
        | 主视图 | Zero123++ 30° 视角 | 正面 |
        | 右视图 | Zero123++ 90° 视角 | 右侧 |
        | 后视图 | Zero123++ 210° 视角 | 背面 |
        | 左视图 | Zero123++ 270° 视角 | 左侧 |
        | 俯视图 | 输入旋转-90° → Zero123++ | 从上往下 |
        | 底视图 | 输入旋转+90° → Zero123++ | 从下往上 |

        ### 技术说明
        - 模型: [Zero123++ v1.2](https://huggingface.co/sudo-ai/zero123plus-v1.2)
        - Zero123++ 固定输出 6 个环绕视图
        - 通过选择合适的视角 + 旋转输入实现工程视图
        - v1.2 改进: 更稳定的视角,FOV 30°

        ### 引用
        ```bibtex
        @misc{shi2023zero123plus,
            title={Zero123++: a Single Image to Consistent Multi-view Diffusion Base Model},
            author={Ruoxi Shi and Hansheng Chen and others},
            year={2023},
            eprint={2310.15110},
            archivePrefix={arXiv},
            primaryClass={cs.CV}
        }
        ```
        """)

        # 绑定事件
        generate_btn.click(
            fn=process_image,
            inputs=[input_image],
            outputs=output_image
        )

    return demo

if __name__ == "__main__":
    # 预加载模型
    print("=" * 50)
    print("Zero123++ 工程六视图生成器")
    print("=" * 50)
    load_model()

    # 启动 demo
    demo = create_demo()
    demo.queue(max_size=5)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )