Update app.py

app.py

@@ -1,53 +1,61 @@
 import gradio as gr
-import torch
-from diffusers import FluxPriorReduxPipeline, FluxPipeline
-from diffusers.utils import load_image
-from huggingface_hub import login
+import requests
 from PIL import Image
-import os
+from io import BytesIO
 
 # 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")
+HF_API_KEY = os.getenv("HF_API_KEY")  # 環境変数からAPIキーを取得
+API_URL_PRIOR = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-Redux-dev"
+API_URL_FLUX = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 
-    # Prior Reduxパイプラインで事前処理
-    pipe_prior_output = pipe_prior_redux(image)
+headers = {
+    "Authorization": f"Bearer {HF_API_KEY}"
+}
 
-    # FLUXパイプラインで画像生成
-    images = pipe(
-        guidance_scale=2.5,
-        num_inference_steps=50,
-        generator=torch.Generator("cpu").manual_seed(0),  # CPU用の乱数生成器を使用
-        **pipe_prior_output,
-    ).images
+def call_hf_api(api_url, payload):
+    response = requests.post(api_url, headers=headers, json=payload)
+    if response.status_code != 200:
+        raise Exception(f"Request failed with status {response.status_code}, {response.text}")
+    return response.json()
 
-    # 生成された画像を返す
-    return images[0]
+def process_image_with_api(image_path):
+    # 入力画像をロード
+    with open(image_path, "rb") as f:
+        image_bytes = f.read()
+    
+    # Prior Reduxモデルで事前処理
+    prior_payload = {"inputs": image_bytes}
+    prior_response = call_hf_api(API_URL_PRIOR, prior_payload)
+    
+    # FLUXモデルで画像生成
+    flux_payload = {
+        "inputs": prior_response,  # Prior Reduxの出力をFLUXモデルに渡す
+        "parameters": {
+            "guidance_scale": 2.5,
+            "num_inference_steps": 50,
+            "seed": 0,  # 再現性のためのシード値
+        }
+    }
+    flux_response = call_hf_api(API_URL_FLUX, flux_payload)
+    
+    # 生成された画像を取得
+    generated_image_url = flux_response.get("generated_image_url")
+    if not generated_image_url:
+        raise Exception("Generated image URL not found in the response.")
+    
+    # URLから画像をダウンロード
+    response = requests.get(generated_image_url)
+    generated_image = Image.open(BytesIO(response.content))
+    
+    return generated_image
 
 # Gradioインターフェースを構築
 def infer(image):
-    result_image = process_image(image)
+    result_image = process_image_with_api(image)
     return result_image
 
 with gr.Blocks() as demo:
-    gr.Markdown("# FLUX Image Generation App")
+    gr.Markdown("# FLUX Image Generation App (via Hugging Face API)")
     
     with gr.Row():
         input_image = gr.Image(type="filepath", label="Input Image")