from openai import OpenAI
import base64
import requests
from dotenv import load_dotenv
import os

load_dotenv()

client = OpenAI(
	base_url="https://ai.gitee.com/v1",
	api_key=os.getenv("api_key"),
	default_headers={"X-Failover-Enabled":"true"},
)

response = client.images.generate(
	prompt="机器人坐在露天草地上，正在画布上画画", # 替换为您的实际提示
	model="flux-1-schnell",
	size="1024x1024",
	extra_body={
		"guidance_scale": 7.5,
		"seed": 42,
	},
)

for i, image_data in enumerate(response.data):
	if image_data.url:
		# 从网址下载
		ext = image_data.url.split('.')[-1].split('?')[0] or "jpg"
		filename = f"flux-1-schnell-output-{i}.{ext}"
		response = requests.get(image_data.url, timeout=30)
		response.raise_for_status()
		with open(filename, "wb") as f:
			f.write(response.content)
		print(f"Downloaded image to {filename}")
	elif image_data.b64_json:
		# 解码 Base64
		image_bytes = base64.b64decode(image_data.b64_json)
		filename = f"flux-1-schnell-output-{i}.jpg"
		with open(filename, "wb") as f:
			f.write(image_bytes)
		print(f"Saved image to {filename}")