import os
import base64
import requests

def text_to_image(prompt):
    url = "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image"

    body = {
      "steps": 30,
      "width": 1024,
      "height": 1024,
      "seed": 0,
      "cfg_scale": 7,
      "samples": 1,
      "style_preset": "enhance",
      "text_prompts": [
        {
          "text": prompt,
          "weight": 1
        }
      ],
    }

    headers = {
      "Accept": "application/json",
      "Content-Type": "application/json",
      "Authorization": "Bearer sk-8n2tgRL6CzV91nQ2nvecVMhedg4GbEexDtUoDnCOWt09Wq2W",
    }

    response = requests.post(
        url,
        headers=headers,
        json=body,
        )

    if response.status_code != 200:
        raise Exception("Non-200 response: " + str(response.text))

    data = response.json()

    for i, image in enumerate(data["artifacts"]):
        if not os.path.exists("output"):
            os.makedirs("output")
            print("成功创建文件夹 'output'")
        image_path = f'./output/txt2img_{image["seed"]}.png'
        with open(image_path, "wb") as f:
            f.write(base64.b64decode(image["base64"]))

    return image_path