import requests
import os
from requests_toolbelt.multipart.encoder import MultipartEncoder
import gradio as gr

def generate_image(prompt, negative_prompt=None, api_key=None, aspect_ratio='1:1', model='sd3', seed=0, output_format='png'):
    api_url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
    #api_key = os.environ.get("STABILITY_AI_API_KEY")  # Retrieve the API key from the environment variable

    fields = {
        'prompt': prompt,
        'aspect_ratio': aspect_ratio,
        'model': model,
        'seed': str(seed),
        'output_format': output_format,
        'mode': 'text-to-image'  # Default mode
    }

    if negative_prompt:
        fields['negative_prompt'] = negative_prompt

    m = MultipartEncoder(fields=fields)

    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': m.content_type,
        'Accept': 'image/*'  # To receive the image directly
    }

    response = requests.post(api_url, data=m, headers=headers)
    print(response.status_code, response.content)
    
    if response.status_code == 200:
        # Assuming the response content is the image in binary format
        output_path = 'generated_image.png'
        with open(output_path, 'wb') as f:
            f.write(response.content)
        return output_path  # Return the path for Gradio to display the image
    else:
        return f"Error: {response.text}"

def wrap_generate_image(prompt, negative_prompt, api_key, model, aspect_ratio):
    return generate_image(prompt, negative_prompt, api_key, aspect_ratio, model)

iface = gr.Interface(
    fn=wrap_generate_image,
    inputs=[
        gr.Textbox(lines=2, label="提示词", placeholder="请输入提示词，例如: 1girl, epic game, miku, city..."),
        gr.Textbox(lines=2, label="负面提示词", placeholder="(worst quality, low quality:1.4)..."),
        gr.Textbox(lines=2, label="API_key", placeholder="填写你的API_KEY..."),
        gr.Radio(choices=['sd3', 'sd3-turbo'], label="SD3模型种类", value='sd3'),
        gr.Dropdown(choices=['1:1', '16:9', '21:9', '2:3', '3:2', '4:5', '5:4', '9:16', '9:21'], label="图片比率", value='1:1')
    ],
    outputs=gr.Image(),
        title="Stable Diffusion 3 在线使用(API)",
    description='''BY BiliBili NYAN9，官网API获取：https://platform.stability.ai/account/keys 。无限账号注册：https://www.snapmail.cc/
    每个账号送25积分，可以免费画五张图左右。后续我会想办法弄个批量注册机。
    支持一下本人吧：https://afdian.net/a/KaggleSD ， 不赞助也没事也可以点进去看看我写过其它的SD项目
    '''
)

iface.launch(share=True)