"""
AI绘画综合示例：支持本地模型与云端API两种模式
功能包含：文本生成图像、图像风格转换、批量生成与结果优化
"""
import os
import time
import requests
import numpy as np
from PIL import Image
from io import BytesIO
import torch
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
# ====================
# 模块1：Stable Diffusion本地模型
# ====================
class StableDiffusionPainter:
def __init__(self, model_name="runwayml/stable-diffusion-v1-5"):
self.pipe = StableDiffusionPipeline.from_pretrained(
model_name,
torch_dtype=torch.float16,
safety_checker=None, # 禁用安全检查以提升生成多样性:ml-citation{ref="2" data="citationList"}
scheduler=EulerAncestralDiscreteScheduler.from_pretrained(model_name, subfolder="scheduler")
)
self.pipe = self.pipe.to("cuda" if torch.cuda.is_available() else "cpu")
def generate(self, prompt, negative_prompt=None, **kwargs):
"""参数说明：
prompt: 生成提示词（中英文皆可）
steps: 迭代步数(20-50)
guidance_scale: 引导系数(7-15)
seed: 随机种子
"""
return self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=kwargs.get('steps', 35),
guidance_scale=kwargs.get('guidance_scale', 7.5),
generator=torch.Generator().manual_seed(kwargs.get('seed', int(time.time())))
).images
# ====================
# 模块2：云端API接口（阿里云示例）
# ====================
class CloudPainter:
def __init__(self, api_key):
self.api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
def generate(self, prompt, style="digital-art", size="1024x1024"):
"""支持风格选项：watercolor/oil-painting/comic/cartoon"""
payload = {
"model": "wanx-v1",
"input": {
"prompt": prompt,
"style": style,
"size": size,
"n": 1
}
}
response = requests.post(self.api_url, json=payload, headers=self.headers)
if response.status_code == 200:
return Image.open(BytesIO(requests.get(response.json()['output']['results']['url']).content))
raise Exception(f"API调用失败: {response.text}")
# ====================
# 模块3：图像后处理
# ====================
class ImageEnhancer:
@staticmethod
def apply_cartoon_effect(image):
"""卡通化处理（参考论文白盒表示方法:ml-citation{ref="1" data="citationList"}）"""
arr = np.array(image)
# 边缘强化
gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 9)
# 颜色量化
quantized = cv2.bilateralFilter(arr, 9, 75, 75)
# 混合处理
cartoon = cv2.bitwise_and(quantized, quantized, mask=edges)
return Image.fromarray(cartoon)
# ====================
# 主程序
# ====================
if __name__ == "__main__":
# 示例1：使用本地模型生成
local_painter = StableDiffusionPainter()
img1 = local_painter.generate(
prompt="Cyberpunk cityscape with neon lights, rain wet streets, 8k detailed",
negative_prompt="blurry, low resolution",
steps=45,
guidance_scale=9.5,
seed=42
)
img1.save("cyberpunk_city.png")
# 示例2：使用阿里云API生成并增强
cloud_painter = CloudPainter(api_key="your_aliyun_key") # 替换真实API密钥:ml-citation{ref="3" data="citationList"}
img2 = cloud_painter.generate("宁静的江南水乡，白墙黛瓦，桃花盛开")
enhancer = ImageEnhancer()
cartoon_img = enhancer.apply_cartoon_effect(img2)
cartoon_img.save("watertown_cartoon.png")
# 典型参数组合参考：
portrait_params = {
"steps": 50, # 人像建议更高步数
"guidance_scale": 12, # 需要精确控制时提高该值
"seed": 12345 # 固定种子可复现结果
}
landscape_params = {
"steps": 35, # 风景可适当降低步数
"guidance_scale": 7.5 # 保持创作自由度
}
# 基础依赖安装（需Python 3.8+）
pip install torch==2.0.1 diffusers transformers accelerate safetensors pillow requests opencv-python
# 验证CUDA可用性（NVIDIA显卡需额外安装CUDA 11.7+）
nvidia-smi # 查看显卡状态
def batch_generate(prompts, painter, output_dir="outputs"):
os.makedirs(output_dir, exist_ok=True)
for i, prompt in enumerate(prompts):
img = painter.generate(prompt)
img.save(f"{output_dir}/result_{i}.png")