调用stable diffusion model的python代码：pipeline自动网络路径下载，内存要求10g左右,模型文件的下载可以国内镜像源

import os
import warnings
import torch
from diffusers import StableDiffusionPipeline

# -----------------------------
# 忽略非关键性警告（可选）
# -----------------------------
warnings.filterwarnings("ignore", category=FutureWarning)
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

# -----------------------------
# 检查 accelerate 是否安装（推荐）
# -----------------------------
try:
    import accelerate
except ImportError:
    print("""
    ⚠️ 注意：未安装 accelerate 库！

    推荐使用 pip 安装以提升加载效率：
    pip install accelerate

    如果你有 GPU，还可以安装 CUDA 版本 PyTorch 以获得更好性能。
    """)
    pass

# -----------------------------
# 配置参数
# -----------------------------
MODEL_NAME = "stabilityai/stable-diffusion-2"  # 可更换为你喜欢的模型
SAVE_DIR = "generated_emotion_dataset"
EMOTIONS = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprise"]
NUM_IMAGES_PER_EMOTION = 100  # 每种情绪生成多少张图片
IMAGE_SIZE = (512, 512)  # 图像尺寸

# 创建文件夹
os.makedirs(SAVE_DIR, exist_ok=True)
for emotion in EMOTIONS:
    os.makedirs(os.path.join(SAVE_DIR, emotion), exist_ok=True)

# -----------------------------
# 加载 Stable Diffusion 模型
# -----------------------------
print("⏳ 正在加载 Stable Diffusion 模型（第一次运行会下载模型）...")

pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    use_safetensors=True,  # 更快更省内存
    safety_checker=None  # 去除 NSFW 检查
)

if torch.cuda.is_available():
    pipe = pipe.to("cuda")


# -----------------------------
# 生成图像函数
# -----------------------------
def generate_images(prompt, save_dir, num_images):
    existing_count = len([f for f in os.listdir(save_dir) if f.endswith(".jpg")])
    remaining = num_images - existing_count

    if remaining <= 0:
        print(f"✅ {prompt} 已生成完毕，跳过。")
        return

    print(f"🎨 正在生成：{prompt}（还需 {remaining} 张）")
    for i in range(existing_count, num_images):
        try:
            image = pipe(prompt).images[0]
            image = image.resize(IMAGE_SIZE)
            image.save(os.path.join(save_dir, f"{i}.jpg"))
            print(f"🖼️ 已生成第 {i + 1}/{num_images} 张图像")
        except Exception as e:
            print(f"❌ 生成失败: {e}")


# -----------------------------
# 主循环生成所有情绪图像
# -----------------------------
for emotion in EMOTIONS:
    prompt = (
        f"A realistic frontal view of an Asian person with a {emotion} expression, "
        f"facing the camera directly, highly detailed, clear expression, close-up, white background"
    )
    emotion_dir = os.path.join(SAVE_DIR, emotion)
    generate_images(prompt, emotion_dir, NUM_IMAGES_PER_EMOTION)

print(f"\n✅ 数据集已生成完成，保存路径：{SAVE_DIR}")