import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch
from PIL import Image
import requests
from io import BytesIO

# デバイスを設定（CUDAが利用可能ならGPUを使用）
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "dalle-mini/dalle-mini"  # 使用するモデルのリポジトリID

# モデルをロード
pipe = DiffusionPipeline.from_pretrained(model_repo_id)
pipe = pipe.to(device)  # モデルを選択したデバイスに移動

# 最大のシード値を定義
MAX_SEED = np.iinfo(np.int32).max

# 画像生成関数
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    # シードをランダム化するオプションの処理
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # シードを指定して画像生成のための乱数生成器を作成
    generator = torch.Generator().manual_seed(seed)

    # プロンプトから画像を生成
    image = pipe(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator
    ).images[0]  # 生成された画像のリストから最初の画像を取得
    
    # 画像を保存
    image.save("generated_image.png")
    
    # 生成された画像とダウンロードリンク、シードを返す
    return image, "generated_image.png", seed

# サンプルプロンプトのリスト
examples = [
    "ジャングルの中の宇宙飛行士、寒色のパレット、 muted colors、詳細、8k",
    "緑の馬に乗った宇宙飛行士",
    "美味しそうなセビーチェチーズケーキスライス",
]

# スタイルシートの設定
css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

# Gradioのインターフェースを構築
with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        # タイトルの表示
        gr.Markdown(f"""
        # テキストから画像への生成器
        """)
        
        # プロンプト入力と実行ボタン
        with gr.Row():
            prompt = gr.Textbox(
                label="プロンプト",
                show_label=False,
                max_lines=1,
                placeholder="プロンプトを入力してください",
                container=False,
            )
            
            run_button = gr.Button("生成", scale=0)  # 画像生成ボタン
        
        result = gr.Image(label="結果", show_label=False)  # 生成された画像を表示するための領域
        download_link = gr.File(label="生成された画像をダウンロード")  # 画像をダウンロードするリンク

        # 詳細設定のアコーディオン
        with gr.Accordion("詳細設定", open=False):
            negative_prompt = gr.Textbox(
                label="ネガティブプロンプト",
                max_lines=1,
                placeholder="ネガティブプロンプトを入力してください",
                visible=False,
            )
            
            seed = gr.Slider(
                label="シード",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="シードをランダム化", value=True)  # シードのランダム化オプション
            
            # 画像の幅と高さを設定するスライダー
            with gr.Row():
                width = gr.Slider(
                    label="幅",
                    minimum=256,
                    maximum=1024,
                    step=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="高さ",
                    minimum=256,
                    maximum=1024,
                    step=32,
                    value=1024,
                )
            
            # ガイダンススケールと推論ステップ数を設定するスライダー
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="ガイダンススケール",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=7.5,
                )
                
                num_inference_steps = gr.Slider(
                    label="推論ステップ数",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=20,
                )
        
        # サンプルプロンプトを提供する機能
        gr.Examples(
            examples=examples,
            inputs=[prompt]
        )

    # ボタンをクリックまたはプロンプトを送信すると画像生成をトリガー
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, download_link, seed]
    )

# Gradioアプリを起動
demo.queue().launch()