Update app.py

app.py

@@ -1,44 +1,41 @@
 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
+from huggingface_hub import login
 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に変更
-
+# Lazy Loadingを実現するため、モデルロードは関数内で行う
 def process_image(image_path):
-    # 入力画像をロード
+    # 入力画像をロードしてリサイズ
     image = Image.open(image_path).convert("RGB")
+    image = image.resize((256, 256))  # サイズを256x256に制限してメモリ節約
 
-    # Prior Reduxパイプラインで事前処理
+    # Prior Reduxパイプラインのロードと処理
+    pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-Redux-dev",
+        use_auth_token=True
+    ).to("cpu")
     pipe_prior_output = pipe_prior_redux(image)
 
-    # FLUXパイプラインで画像生成
+    # FLUXパイプラインのロードと処理
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        use_auth_token=True,
+        text_encoder=None,
+        text_encoder_2=None
+    ).to("cpu")
     images = pipe(
         guidance_scale=2.5,
-        num_inference_steps=50,
-        generator=torch.Generator("cpu").manual_seed(0),  # CPU用の乱数生成器を使用
+        num_inference_steps=25,  # 推論ステップを減らしてメモリ節約
+        generator=torch.Generator("cpu").manual_seed(0),  # 再現性のためのシード値
         **pipe_prior_output,
     ).images
 
-    # 生成された画像を返す
+    # 結果画像を返す
     return images[0]
 
 # Gradioインターフェースを構築
@@ -47,8 +44,8 @@ def infer(image):
     return result_image
 
 with gr.Blocks() as demo:
-    gr.Markdown("# FLUX Image Generation App")
-    
+    gr.Markdown("# FLUX Image Generation App (Optimized for CPU)")
+
     with gr.Row():
         input_image = gr.Image(type="filepath", label="Input Image")
         output_image = gr.Image(type="pil", label="Generated Image")