Update app.py

app.py

@@ -1,64 +1,53 @@
 import gradio as gr
-import requests
+import torch
+from diffusers import FluxPriorReduxPipeline, FluxPipeline
+from diffusers.utils import load_image
+from huggingface_hub import login
 from PIL import Image
-from io import BytesIO
 import os
 
 # Hugging Face HubのAPIキーを設定
-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"
-
-headers = {
-    "Authorization": f"Bearer {HF_API_KEY}"
-}
+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")
 
-def call_hf_api(api_url, image_bytes, parameters=None):
-    # 画像データを multipart/form-data として送信
-    files = {
-        "file": ("image.png", image_bytes, "image/png"),
-    }
-    data = {"parameters": parameters} if parameters else {}
+    # Prior Reduxパイプラインで事前処理
+    pipe_prior_output = pipe_prior_redux(image)
 
-    response = requests.post(api_url, headers=headers, files=files, data=data)
-    if response.status_code != 200:
-        raise Exception(f"Request failed with status {response.status_code}, {response.text}")
-    return response.json()
+    # FLUXパイプラインで画像生成
+    images = pipe(
+        guidance_scale=2.5,
+        num_inference_steps=50,
+        generator=torch.Generator("cpu").manual_seed(0),  # CPU用の乱数生成器を使用
+        **pipe_prior_output,
+    ).images
 
-def process_image_with_api(image_path):
-    # 入力画像をロード
-    with open(image_path, "rb") as f:
-        image_bytes = f.read()
-    
-    # Prior Reduxモデルで事前処理
-    prior_response = call_hf_api(API_URL_PRIOR, image_bytes)
-    
-    # FLUXモデルで画像生成
-    flux_payload = {
-        "guidance_scale": 2.5,
-        "num_inference_steps": 50,
-        "seed": 0,  # 再現性のためのシード値
-    }
-    flux_response = call_hf_api(API_URL_FLUX, image_bytes, parameters=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
+    # 生成された画像を返す
+    return images[0]
 
 # Gradioインターフェースを構築
 def infer(image):
-    result_image = process_image_with_api(image)
+    result_image = process_image(image)
     return result_image
 
 with gr.Blocks() as demo:
-    gr.Markdown("# FLUX Image Generation App (via Hugging Face API)")
+    gr.Markdown("# FLUX Image Generation App")
     
     with gr.Row():
         input_image = gr.Image(type="filepath", label="Input Image")