import gradio as gr
import requests
import io
import os
import json
import random
from PIL import Image

def query(model_name, prompt, is_negative, steps, cfg_scale, seed, strength, width, height):
    API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
    headers = {"Authorization": f"Bearer {os.getenv('token')}"}
    
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,
            "height": height
        }
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    
    if response.status_code != 200:
        return f"Error {response.status_code}: {response.text}"
    
    try:
        image = Image.open(io.BytesIO(response.content))
        return image
    except Exception:
        return f"Invalid response: {response.text}"

demo = gr.Interface(
    fn=query,
    inputs=[
        gr.Textbox(value="strangerzonehf/Flux-Animeo-v1-LoRA", label="モデル名"),
        gr.Textbox(label="プロンプト", placeholder="生成する画像の説明を入力"),
        gr.Checkbox(label="ネガティブプロンプトを使用"),
        gr.Slider(1, 100, step=1, value=25, label="ステップ数"),
        gr.Slider(1, 20, step=0.5, value=7, label="CFGスケール"),
        gr.Number(-1, label="シード (-1でランダム)"),
        gr.Slider(0.0, 1.0, step=0.1, value=0.7, label="変換強度"),
        gr.Slider(256, 1024, step=64, value=512, label="幅"),
        gr.Slider(256, 1024, step=64, value=512, label="高さ")
    ],
    outputs=gr.Image(label="生成画像"),
    title="Flux-Animeo 画像生成",
    description="テキストプロンプトとオプションを設定して画像を生成します。"
)

demo.launch()