import gradio as gr
import torch
from diffusers import FluxPriorReduxPipeline, FluxPipeline
from diffusers.utils import load_image
from huggingface_hub import login
from PIL import Image
import os

# Hugging Face HubのAPIキーを設定
login(os.getenv("HF_API_KEY"))

# Prior Reduxパイプラインをロード (CPU対応)
pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Redux-dev",
    use_auth_token=True  # APIキーを使用して認証
).to("cpu")  # CPUに変更

# メインのFLUXパイプラインをロード (CPU対応)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    use_auth_token=True,  # APIキーを使用して認証
    text_encoder=None,
    text_encoder_2=None
).to("cpu")  # CPUに変更

def process_image(image_path):
    # 入力画像をロード
    image = Image.open(image_path).convert("RGB")

    # Prior Reduxパイプラインで事前処理
    pipe_prior_output = pipe_prior_redux(image)

    # FLUXパイプラインで画像生成
    images = pipe(
        guidance_scale=2.5,
        num_inference_steps=50,
        generator=torch.Generator("cpu").manual_seed(0),  # CPU用の乱数生成器を使用
        **pipe_prior_output,
    ).images

    # 生成された画像を返す
    return images[0]

# Gradioインターフェースを構築
def infer(image):
    result_image = process_image(image)
    return result_image

with gr.Blocks() as demo:
    gr.Markdown("# FLUX Image Generation App")
    
    with gr.Row():
        input_image = gr.Image(type="filepath", label="Input Image")
        output_image = gr.Image(type="pil", label="Generated Image")

    submit_button = gr.Button("Generate")

    submit_button.click(fn=infer, inputs=[input_image], outputs=[output_image])

demo.launch()