【Streamlit】 BFLのFLUX Pro 1.1-ultraで画像生成する

import streamlit as st
import requests
from PIL import Image
from io import BytesIO
import os
import time

# 環境変数からAPIキーを取得
API_KEY = os.getenv('BFL_API_KEY')

# APIエンドポイント
REQUEST_URL = 'https://api.us1.bfl.ai/v1/flux-pro-1.1-ultra'
RESULT_URL = 'https://api.us1.bfl.ai/v1/get_result'

# Streamlit UI
st.title('Image Generation with FLUX-DEV Model')

prompt = st.text_input('Enter your prompt:')
width = st.number_input('Width of the image:', min_value=64, max_value=2048, value=512, step=64)
height = st.number_input('Height of the image:', min_value=64, max_value=2048, value=512, step=64)

if st.button('Generate Image'):
    if prompt:
        with st.spinner('Generating image...'):
            # リクエストペイロードを作成
            payload = {
                'prompt': prompt,
                'width': width,
                'height': height,
                'output_format': 'png',
                'raw': False,
            }
            headers = {
                'accept': 'application/json',
                'x-key': API_KEY,
                'Content-Type': 'application/json'
            }

            # 画像生成リクエストを送信
            response = requests.post(REQUEST_URL, json=payload, headers=headers)
            if response.status_code == 200:
                request_id = response.json().get('id')
                if request_id:
                    # 結果を取得するためのポーリング処理
                    while True:
                        time.sleep(0.5)
                        result_response = requests.get(RESULT_URL, headers=headers, params={'id': request_id})
                        if result_response.status_code == 200:
                            result_data = result_response.json()
                            if result_data.get('status') == 'Ready':
                                image_url = result_data['result'].get('sample')
                                if image_url:
                                    # 画像を取得し表示
                                    image_response = requests.get(image_url)
                                    if image_response.status_code == 200:
                                        image = Image.open(BytesIO(image_response.content))
                                        st.image(image, caption='Generated Image', use_column_width=True)
                                    else:
                                        st.error('Failed to retrieve the generated image.')
                                else:
                                    st.error('Image URL not found in the response.')
                                break
                            elif result_data.get('status') == 'Failed':
                                st.error('Image generation failed.')
                                break
                        else:
                            st.error('Error while fetching the result.')
                else:
                    st.error('Request ID not found in the response.')
            else:
                st.error('Error while submitting the generation request.')
    else:
        st.warning('Please enter a prompt.')

# Two people watching fireworks at Disneyland.