前言

人工智能技术的发展，使得图像生成成为了创意工作者和技术爱好者的重要工具之一。Civitai 是一个强大的 AI 图像生成平台，通过简洁易用的界面和强大的生成能力，允许用户轻松创建高质量的图像。本文介绍获取生成图片和提示信息实现一个简化web

源码地址

示例

读取配置

首先，我们需要配置脚本以确保它能正确地运行。我们使用 configparser 模块来读取配置文件 config.ini。该配置文件包含了一些关键参数，如输出目录、API 密钥等。

import os
import json
import requests
import configparser

# Configuration
config = configparser.ConfigParser()
config.read("config.ini")
output_dir = config["DEFAULT"]["OutputDir"]
nsfw_level = config["DEFAULT"]["NSFWLevel"]
api_key = config["DEFAULT"]["ApiKey"]
local_images = config["DEFAULT"]["LocalImages"].lower() == 'true'
cursor = config["DEFAULT"]["Cursor"]

保存元数据

接下来，我们定义一个函数 save_metadata 来保存图片的元数据。该函数会将元数据保存为 JSON 格式的文件。

def save_metadata(item, output_dir):
    item_id = item["id"]
    json_filename = os.path.join(output_dir, f"{item_id}.json")
    if os.path.isfile(json_filename):
        print(f"File {json_filename} already exists.")
    else:
        print(f"Saving {json_filename}.")
        with open(json_filename, "w") as f:
            json.dump(item, f, indent=1)

保存图片

为了保存图片，我们定义了另一个函数 save_image。该函数会检查是否需要本地保存图片，如果是，则下载图片并保存在指定目录下。

def save_image(item, output_dir):
    if item["meta"] and item["url"] and local_images:
        img_filename = os.path.join(output_dir, f"{item['id']}.jpg")
        image_response = requests.get(item["url"])
        if image_response.status_code == 200:
            with open(img_filename, "wb") as f:
                f.write(image_response.content)
            print("Image saved successfully!")
        else:
            print("Error saving image:", image_response.status_code)

获取图片

定义 fetch_images 函数来发送 GET 请求以获取图片数据。

def fetch_images(url):
    print(f"Fetching URL: {url}")
    response = requests.get(url)
    return response

运行

最后，我们编写主逻辑以循环获取图片数据并保存元数据和图片。我们将分页处理数据，并根据 API 返回的 nextCursor 来处理分页。

# Initial URL
url = f"https://civitai.com/api/v1/images?token={api_key}&limit=100&sort=Most%20Reactions&nsfw={nsfw_level}&cursor={cursor}"

while url:
    response = fetch_images(url)

    if response.status_code == 200:
        data = response.json()
        next_cursor = data["metadata"].get("nextCursor")
        
        if next_cursor:
            # Convert nextCursor to directory-friendly format
            cursor_dir_name = next_cursor.replace('|', '_')
            # Create a new directory named after converted nextCursor
            page_output_dir = os.path.join(output_dir, cursor_dir_name)
            os.makedirs(page_output_dir, exist_ok=True)

            for item in data["items"]:
                save_metadata(item, page_output_dir)
                save_image(item, page_output_dir)
            
            # Update URL to nextPage if it exists
            url = data["metadata"].get("nextPage", None)
        else:
            break
    else:
        print("Error:", response.status_code)
        break