import os
import torch
import gradio as gr
from diffusers import DiffusionPipeline

# ===== モデル設定 =====
MODEL_ID = os.environ.get("MODEL_ID", "prompthero/openjourney")

# 無料CPU前提：float32が安定（bfloat16/float16はCPUだと非推奨）
torch_dtype = torch.float32
device = "cpu"  # GPUを使うなら "cuda" に変更（SpacesのHWもGPUへ）

# パイプラインをロード
# 注意: from_pretrainedの引数はCPU/FP32に合わせて簡素化
pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
    safety_checker=None  # 必要なら独自にNSFWフィルタを実装
)
pipe = pipe.to(device)

# 推論関数
def generate_image(prompt, steps, guidance, seed, width, height):
    # 乱数シード（再現性）
    generator = None
    if seed is not None and seed != "":
        try:
            generator = torch.Generator(device=device).manual_seed(int(seed))
        except Exception:
            generator = None

    # CPUではサイズを抑えると速い（例: 512x512）
    result = pipe(
        prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(width),
        height=int(height),
        generator=generator
    )
    image = result.images[0]
    return image

# Gradio UI
with gr.Blocks(theme="soft") as demo:
    gr.Markdown(
        "# 🎨 OpenJourney 画像生成（CPU/Free）\n"
        "無料CPUで動作するため、生成には時間がかかります。サイズとステップを小さめにすると速くなります。"
    )

    with gr.Row():
        prompt = gr.Textbox(
            label="プロンプト",
            value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
        )

    with gr.Row():
        steps = gr.Slider(10, 50, value=25, step=1, label="num_inference_steps（多いほど高品質・遅い）")
        guidance = gr.Slider(1.0, 12.0, value=7.5, step=0.1, label="guidance_scale（プロンプト忠実度）")

    with gr.Row():
        width = gr.Dropdown(choices=["384","448","512","576","640"], value="512", label="幅（px）")
        height = gr.Dropdown(choices=["384","448","512","576","640"], value="512", label="高さ（px）")
        seed = gr.Textbox(value="", label="seed（空ならランダム）")

    generate_btn = gr.Button("生成")
    output = gr.Image(label="出力画像", type="pil")

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, steps, guidance, seed, width, height],
        outputs=[output]
    )

if __name__ == "__main__":
    demo.launch()