AI绘图与工业设计：从概念到产品_03

import torch
import numpy as np
import cv2
from PIL import Image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
import gradio as gr
import os
from datetime import datetime

# 设置中文字体支持
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 模型加载与初始化
class IndustrialDesignGenerator:
    def __init__(self, model_name="runwayml/stable-diffusion-v1-5", device="cuda"):
        self.device = device if torch.cuda.is_available() else "cpu"
        print(f"使用设备: {self.device}")

        # 加载ControlNet模型
        self.controlnet_canny = ControlNetModel.from_pretrained(
            "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16
        )
        self.controlnet_depth = ControlNetModel.from_pretrained(
            "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
        )

        # 加载主模型
        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
            model_name,
            controlnet=[self.controlnet_canny, self.controlnet_depth],
            torch_dtype=torch.float16,
        )
        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
        self.pipe.to(self.device)

        # 启用安全检查器
        self.pipe.safety_checker = lambda images, clip_input: (images, False)

        # 预设设计风格提示词
        self.design_styles = {
            "极简主义": "minimalist design, clean lines, simple form, monochromatic, high quality materials",
            "工业风": "industrial design, raw materials, exposed structure, functional, utilitarian",
            "未来主义": "futuristic design, sleek, metallic, advanced technology, innovative features",
            "有机设计": "organic design, natural shapes, flowing forms, biomimetic, sustainable materials",
            "复古风格": "retro design, nostalgic elements, classic proportions, vintage details"
        }

        # 预设产品类型提示词模板
        self.product_templates = {
            "智能手表": "smartwatch, circular display, stainless steel case, silicone strap, minimal interface, {style}, high resolution render, product photography, studio lighting",
            "蓝牙耳机": "wireless earbuds, compact charging case, ergonomic design, {style}, high resolution render, product photography, white background",
            "咖啡机": "coffee machine, compact size, user friendly interface, {style}, stainless steel and glass materials, high resolution render, kitchen environment",
            "办公椅": "office chair, ergonomic design, adjustable height, {style}, comfortable upholstery, high resolution render, office environment",
            "灯具": "lighting fixture, ambient lighting, {style}, high quality materials, energy efficient, high resolution render, living room setting"
        }

    def preprocess_image(self, image, low_threshold=100, high_threshold=200):
        """预处理图像生成边缘图"""
        image = np.array(image)
        image = cv2.Canny(image, low_threshold, high_threshold)
        image = image[:, :, None]
        image = np.concatenate([image, image, image], axis=2)
        return Image.fromarray(image)

    def generate_design(self,
                       product_type,
                       design_style,
                       additional_prompt="",
                       control_image=None,
                       depth_image=None,
                       num_inference_steps=30,
                       guidance_scale=7.5,
                       seed=None):
        """
        生成工业设计图像
        :param product_type: 产品类型
        :param design_style: 设计风格
        :param additional_prompt: 额外提示词
        :param control_image: 控制图像（草图）
        :param depth_image: 深度图像
        :param num_inference_steps: 推理步数
        :param guidance_scale: 引导尺度
        :param seed: 随机种子
        :return: 生成的图像
        """
        # 设置随机种子
        if seed is None:
            seed = torch.randint(0, 1000000, (1,)).item()
        generator = torch.manual_seed(seed)

        # 构建提示词
        if product_type in self.product_templates:
            base_prompt = self.product_templates[product_type]
        else:
            base_prompt = f"{product_type}, {design_style}, high resolution render, product photography"

        style_prompt = self.design_styles.get(design_style, design_style)
        prompt = base_prompt.replace("{style}", style_prompt) + ", " + additional_prompt
        prompt = prompt.strip().rstrip(',')

        # 构建负面提示词
        negative_prompt = "low quality, blurry, distorted, ugly, bad proportions, poorly drawn, extra limbs, missing parts"

        # 预处理控制图像
        if control_image is not None:
            control_image = self.preprocess_image(control_image)
        else:
            # 如果没有控制图像，创建空白图像
            control_image = Image.new('RGB', (512, 512), color='white')
            control_image = self.preprocess_image(control_image)

        # 如果没有深度图像，使用控制图像
        if depth_image is None:
            depth_image = control_image

        # 生成图像
        images = self.pipe(
            prompt,
            image=[control_image, depth_image],
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator,
            controlnet_conditioning_scale=[1.0, 0.8]
        ).images

        return images[0], seed, prompt

    def save_design(self, image, product_type, design_style):
        """保存生成的设计图像"""
        # 创建保存目录
        save_dir = os.path.join("generated_designs", f"{product_type}_{design_style}")
        os.makedirs(save_dir, exist_ok=True)

        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.join(save_dir, f"design_{timestamp}.png")

        # 保存图像
        image.save(filename)
        return filename

# Gradio界面
def create_interface():
    designer = IndustrialDesignGenerator()

    with gr.Blocks(title="工业设计AI助手") as demo:
        gr.Markdown("# 工业设计AI助手")
        gr.Markdown("通过AI快速生成产品设计方案，支持草图控制和风格调整")

        with gr.Row():
            with gr.Column(scale=1):
                product_type = gr.Dropdown(
                    choices=list(designer.product_templates.keys()),
                    label="产品类型",
                    value="智能手表"
                )
                design_style = gr.Dropdown(
                    choices=list(designer.design_styles.keys()),
                    label="设计风格",
                    value="极简主义"
                )
                additional_prompt = gr.Textbox(
                    label="额外提示词",
                    placeholder="输入材质、颜色、功能等额外要求...",
                    value=""
                )
                control_image = gr.Image(
                    label="草图控制图像",
                    type="pil",
                    height=300
                )
                with gr.Accordion("高级设置", open=False):
                    num_inference_steps = gr.Slider(
                        label="推理步数",
                        minimum=10,
                        maximum=100,
                        value=30,
                        step=1
                    )
                    guidance_scale = gr.Slider(
                        label="引导尺度",
                        minimum=1,
                        maximum=20,
                        value=7.5,
                        step=0.5
                    )
                    seed = gr.Number(
                        label="随机种子",
                        value=None,
                        precision=0
                    )
                generate_btn = gr.Button("生成设计", variant="primary")

            with gr.Column(scale=1):
                result_image = gr.Image(label="生成的设计", height=512)
                generated_prompt = gr.Textbox(label="生成的提示词", lines=3)
                seed_display = gr.Textbox(label="使用的种子", interactive=False)
                save_btn = gr.Button("保存设计")
                save_result = gr.Textbox(label="保存结果", interactive=False)

        # 生成按钮事件
        def generate_design_event(*args):
            product_type, design_style, additional_prompt, control_image, num_steps, guidance, seed = args
            image, used_seed, prompt = designer.generate_design(
                product_type=product_type,
                design_style=design_style,
                additional_prompt=additional_prompt,
                control_image=control_image,
                num_inference_steps=num_steps,
                guidance_scale=guidance,
                seed=seed if seed else None
            )
            return image, prompt, str(used_seed)

        generate_btn.click(
            fn=generate_design_event,
            inputs=[product_type, design_style, additional_prompt, control_image, num_inference_steps, guidance_scale, seed],
            outputs=[result_image, generated_prompt, seed_display]
        )

        # 保存按钮事件
        def save_design_event(image, product_type, design_style):
            if image is None:
                return "请先生成设计"
            try:
                filename = designer.save_design(image, product_type, design_style)
                return f"设计已保存至: {filename}"
            except Exception as e:
                return f"保存失败: {str(e)}"

        save_btn.click(
            fn=save_design_event,
            inputs=[result_image, product_type, design_style],
            outputs=[save_result]
        )

    return demo

# 主函数
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)