2026 AI 图像与视频生成全景：从 Flux 到 Sora 2，模型架构、开源生态与实战指南

"""
实战 1：使用 Diffusers + Flux 生成高质量图像
支持提示词优化、多张批量生成、自动保存
"""

import torch
from diffusers import FluxPipeline
from PIL import Image
import os
from datetime import datetime

class FluxImageGenerator:
    """Flux 图像生成器，支持批量生成和参数控制"""
    
    def __init__(
        self,
        model_id: str = "black-forest-labs/FLUX.1-dev",
        device: str = "cuda",
        torch_dtype: torch.dtype = torch.bfloat16,
    ):
        self.device = device
        self.pipe = FluxPipeline.from_pretrained(
            model_id,
            torch_dtype=torch_dtype,
        ).to(device)
        # 启用 xformers 显存优化
        self.pipe.enable_xformers_memory_efficient_attention()
        
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        num_images: int = 1,
        guidance_scale: float = 3.5,
        num_inference_steps: int = 28,
        seed: int = None,
        output_dir: str = "outputs",
    ) -> list[Image.Image]:
        """生成图像并保存到指定目录"""
        os.makedirs(output_dir, exist_ok=True)
        
        # 设置随机种子（可选）
        generator = None
        if seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        
        # 批量生成
        images = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
        ).images
        
        # 保存图像
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        saved_paths = []
        for i, img in enumerate(images):
            filename = f"{timestamp}_img{i}.png"
            filepath = os.path.join(output_dir, filename)
            img.save(filepath)
            saved_paths.append(filepath)
            print(f"✅ 已保存: {filepath}")
        
        return images, saved_paths
    
    def enhance_prompt(self, base_prompt: str) -> str:
        """增强提示词，添加质量关键词"""
        quality_tags = [
            "masterpiece", "best quality", "ultra-detailed",
            "8k resolution", "photorealistic"
        ]
        return f"{base_prompt}, {', '.join(quality_tags)}"

# === 使用示例 ===
if __name__ == "__main__":
    # 初始化生成器
    generator = FluxImageGenerator(
        model_id="black-forest-labs/FLUX.1-schnell",  # 快速版
        device="cuda"
    )
    
    # 生成单张图片
    prompt = "A majestic dragon flying over a medieval castle at sunset, cinematic lighting"
    images, paths = generator.generate(
        prompt=generator.enhance_prompt(prompt),
        num_images=1,
        guidance_scale=3.5,
        num_inference_steps=4,  # schnell 只需要 4 步
        seed=42,
    )
    
    # 批量生成不同种子
    prompts = [
        "Cyberpunk city street at night with neon signs and rain",
        "A cozy coffee shop interior with warm lighting and books",
        "Astronaut floating in space with Earth in background",
    ]
    
    for i, p in enumerate(prompts):
        generator.generate(
            prompt=generator.enhance_prompt(p),
            num_images=1,
            seed=100 + i,
            output_dir=f"outputs/batch_{i}"
        )

"""
实战 2：通过 ComfyUI WebSocket API 自动化图像生成
构建自定义工作流并批量执行
"""

import json
import websocket
import uuid
import urllib.request
import urllib.parse
import time

COMFYUI_URL = "127.0.0.1:8188"
CLIENT_ID = str(uuid.uuid4())

class ComfyUIAutomation:
    """ComfyUI 自动化客户端"""
    
    def __init__(self, server_address: str = COMFYUI_URL):
        self.server_address = server_address
        self.client_id = CLIENT_ID
        
    def build_flux_workflow(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 1024,
        height: int = 1024,
        steps: int = 20,
        cfg: float = 3.5,
        seed: int = None,
    ) -> dict:
        """构建 Flux 文生图工作流"""
        if seed is None:
            seed = int(time.time() * 1000) % (2**32)
            
        return {
            "4": {  # CLIPTextEncode (positive)
                "class_type": "CLIPTextEncodeFlux",
                "inputs": {
                    "clip": ["11", 0],
                    "prompt": prompt,
                    "lora_strength": 1.0,
                }
            },
            "5": {  # CLIPTextEncode (negative)
                "class_type": "CLIPTextEncode",
                "inputs": {
                    "clip": ["11", 0],
                    "text": negative_prompt,
                }
            },
            "6": {  # EmptyLatentImage
                "class_type": "EmptyLatentImage",
                "inputs": {
                    "width": width,
                    "height": height,
                    "batch_size": 1,
                }
            },
            "8": {  # VAEDecode
                "class_type": "VAEDecode",
                "inputs": {
                    "samples": ["13", 0],
                    "vae": ["10", 0],
                }
            },
            "9": {  # SaveImage
                "class_type": "SaveImage",
                "inputs": {
                    "images": ["8", 0],
                    "filename_prefix": "comfy_auto",
                }
            },
            "10": {  # VAELoader (Flux AE)
                "class_type": "VAELoader",
                "inputs": {
                    "vae_name": "ae.safetensors",
                }
            },
            "11": {  # DualCLIPLoader
                "class_type": "DualCLIPLoader",
                "inputs": {
                    "clip_name1": "clip_l.safetensors",
                    "clip_name2": "t5xxl_fp16.safetensors",
                    "type": "flux",
                }
            },
            "12": {  # UNETLoader (Flux)
                "class_type": "UNETLoader",
                "inputs": {
                    "unet_name": "flux1-dev.safetensors",
                    "weight_dtype": "default",
                }
            },
            "13": {  # KSampler
                "class_type": "KSampler",
                "inputs": {
                    "model": ["12", 0],
                    "positive": ["4", 0],
                    "negative": ["5", 0],
                    "latent_image": ["6", 0],
                    "seed": seed,
                    "steps": steps,
                    "cfg": cfg,
                    "sampler_name": "euler",
                    "scheduler": "normal",
                    "denoise": 1.0,
                }
            },
        }
    
    def queue_prompt(self, workflow: dict) -> str:
        """将工作流提交到 ComfyUI 队列"""
        data = json.dumps({
            "prompt": workflow,
            "client_id": self.client_id,
        }).encode("utf-8")
        
        req = urllib.request.Request(
            f"http://{self.server_address}/prompt",
            data=data,
            headers={"Content-Type": "application/json"},
        )
        response = urllib.request.urlopen(req)
        result = json.loads(response.read())
        return result["prompt_id"]
    
    def get_images(self, prompt_id: str) -> list:
        """获取生成结果"""
        req = urllib.request.Request(
            f"http://{self.server_address}/history/{prompt_id}"
        )
        response = urllib.request.urlopen(req)
        history = json.loads(response.read())
        
        outputs = history[prompt_id]["outputs"]
        images = []
        for node_output in outputs.values():
            if "images" in node_output:
                for img in node_output["images"]:
                    filename = img["filename"]
                    images.append(filename)
        return images
    
    def generate_batch(
        self,
        prompts: list[str],
        **workflow_kwargs,
    ) -> list[list[str]]:
        """批量生成，返回每张图片的文件名列表"""
        all_results = []
        for i, prompt in enumerate(prompts):
            print(f"🎨 生成第 {i+1}/{len(prompts)} 张...")
            workflow = self.build_flux_workflow(prompt=prompt, **workflow_kwargs)
            prompt_id = self.queue_prompt(workflow)
            # 等待生成完成
            time.sleep(15)  # 实际应用中应监听 WebSocket 事件
            images = self.get_images(prompt_id)
            all_results.append(images)
            print(f"✅ 完成: {images}")
        return all_results

# === 使用示例 ===
if __name__ == "__main__":
    comfy = ComfyUIAutomation()
    
    prompts = [
        "A futuristic Tokyo street at night, neon lights, rain reflections",
        "A steampunk airship flying over Victorian London",
        "An underwater coral reef with tropical fish, sunlight rays",
    ]
    
    results = comfy.generate_batch(
        prompts=prompts,
        width=1024,
        height=1024,
        steps=25,
        cfg=3.5,
    )
    
    print(f"\n📊 批量生成完成，共生成 {sum(len(r) for r in results)} 张图片")