image-gen-mcp / app.py
baxin's picture
Update app.py
bf6edf9 verified
# app.py
import gradio as gr
import together
import os
import base64
from PIL import Image
from io import BytesIO
import requests # ▼▼▼ 修正箇所: requestsライブラリをインポート ▼▼▼
# --- 1. 設定: APIキーの読み込み ---
api_key = os.environ.get("TOGETHER_API_KEY")
if api_key:
together_client = together.Together(api_key=api_key)
else:
together_client = None
# 使用する画像生成モデル
IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell-Free"
# --- 2. 中核機能: 画像生成関数 ---
def generate_image_from_prompt(prompt: str, history: list):
if not together_client:
raise gr.Error("TogetherAIのAPIクライアントが設定されていません。環境変数 TOGETHER_API_KEY を確認してください。")
if not prompt or not prompt.strip():
raise gr.Error("画像を生成するには、プロンプトを入力してください。")
print(f"プロンプトを元に画像を生成中: '{prompt}'")
try:
response = together_client.images.generate(
prompt=prompt,
model=IMAGE_MODEL,
steps=4,
n=4,
height=1024,
width=1024
)
generated_images = []
# ▼▼▼ 修正箇所: URLから画像をダウンロードするロジックに変更 ▼▼▼
for image_data_item in response.data:
# APIの応答に 'url' が含まれているかチェック
if hasattr(image_data_item, 'url') and image_data_item.url:
try:
image_url = image_data_item.url
print(f"URLから画像をダウンロード中: {image_url}")
# URLにアクセスして画像データを取得
image_response = requests.get(image_url)
image_response.raise_for_status() # エラーがあれば例外を発生させる
# ダウンロードしたバイトデータからPILイメージを作成
image = Image.open(BytesIO(image_response.content))
generated_images.append((image, prompt))
except requests.exceptions.RequestException as e:
print(f"画像のダウンロードに失敗しました: {e}")
continue # この画像の処理をスキップして次に進む
# (参考) b64_jsonが返ってくるモデルのためのフォールバック処理
elif image_data_item.b64_json:
image_b64 = image_data_item.b64_json
image_data = base64.b64decode(image_b64)
image = Image.open(BytesIO(image_data))
generated_images.append((image, prompt))
if not generated_images:
raise gr.Error("画像の生成に失敗しました。APIから有効な画像データやURLが返されませんでした。")
new_history = generated_images + history
print(f"{len(generated_images)}枚の画像の生成が完了しました。")
return new_history, new_history
except Exception as e:
if isinstance(e, gr.Error):
raise e
print(f"予期せぬエラーが発生しました: {e}")
raise gr.Error(f"画像の生成に失敗しました。エラー詳細: {e}")
# --- 3. Gradio UIの構築 (変更なし) ---
with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 800px; margin: auto;}") as app:
gr.Markdown(
"""
# 🖼️ TogetherAI 画像生成アプリ (Gradio版)
プロンプトを入力して、TogetherAIのAPIで画像を生成します。
最新の画像がギャラリーの先頭に表示されます。
"""
)
gr.Markdown("---")
history_state = gr.State([])
with gr.Column():
prompt_input = gr.Textbox(
label="画像生成プロンプト",
placeholder="例: 月面にいる、写真のようにリアルな猫の宇宙飛行士",
lines=4,
interactive=True,
)
if not together_client:
gr.Warning("環境変数 `TOGETHER_API_KEY` が見つかりません。画像生成機能を利用するにはAPIキーを設定してください。")
generate_btn = gr.Button("画像生成 ✨ (4枚)", variant="primary", interactive=(together_client is not None))
gr.Markdown("---")
gr.Markdown("## 生成された画像")
image_gallery = gr.Gallery(label="生成結果", show_label=False, columns=2, height="auto", object_fit="contain", value=None, preview=True)
generate_btn.click(fn=generate_image_from_prompt, inputs=[prompt_input, history_state], outputs=[image_gallery, history_state], show_progress="full")
gr.Examples(
examples=[
"A high-resolution photo of a futuristic city with flying cars",
"An oil painting of a serene landscape with a river and mountains",
"A cute, fluffy alien creature exploring a vibrant, magical forest",
"Logo for a coffee shop named 'The Daily Grind', minimalist, modern"
],
inputs=prompt_input
)
if __name__ == "__main__":
app.launch(mcp_server=True)