Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from fastapi import FastAPI, HTTPException, Body
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -15,11 +15,13 @@ import tempfile
 import time
 import os
 import logging
-import numpy as np  # 追加: 画像処理の最適化用
-import threading  # 追加: 並列処理のため
-import queue  # 追加: WebDriverプール用
-from concurrent.futures import ThreadPoolExecutor  # 追加: 並列処理用
-from huggingface_hub import hf_hub_download
+import numpy as np
+import threading
+import queue
+import uuid
+from datetime import datetime
+from concurrent.futures import ThreadPoolExecutor
+from huggingface_hub import hf_hub_download, upload_file, login
 
 # 正しいGemini関連のインポート
 import google.generativeai as genai
@@ -28,6 +30,77 @@ import google.generativeai as genai
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# --- HuggingFace Hub アップロード機能 ---
+class HuggingFaceUploader:
+    """HuggingFace Hubへ画像をアップロードする機能を提供するクラス"""
+    def __init__(self):
+        self.repo_id = os.environ.get("HF_REPO_ID", "tomo2chin2/SUPER_TENSAI_JIN")
+        self.token = os.environ.get("HF_TOKEN", None)
+        if self.token:
+            try:
+                login(token=self.token)
+                logger.info(f"HuggingFace Hubにログインしました。リポジトリ: {self.repo_id}")
+            except Exception as e:
+                logger.error(f"HuggingFace Hubへのログインに失敗: {e}")
+        else:
+            logger.warning("HF_TOKEN環境変数が設定されていません。アップロード機能は制限されます。")
+    
+    def upload_image(self, image, prefix="generated"):
+        """
+        PIL Imageをアップロードし、アクセス可能なURLを返す
+        
+        Args:
+            image: PIL.Image - アップロードする画像
+            prefix: str - ファイル名のプレフィックス
+            
+        Returns:
+            str - アップロードされた画像のURL
+        """
+        try:
+            if not self.token:
+                logger.error("HF_TOKENが設定されていないため、アップロードできません")
+                return None
+                
+            # ユニークなファイル名を生成
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            unique_id = str(uuid.uuid4())[:8]
+            filename = f"{prefix}_{timestamp}_{unique_id}.png"
+            path_in_repo = f"images/{filename}"
+            
+            # 一時ファイルに保存
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+                tmp_path = tmp_file.name
+                image.save(tmp_path, format="PNG")
+            
+            logger.info(f"画像を一時ファイルに保存: {tmp_path}")
+            
+            # HuggingFaceにアップロード
+            logger.info(f"HuggingFace Hubにアップロード中: {path_in_repo}")
+            upload_info = upload_file(
+                path_or_fileobj=tmp_path,
+                path_in_repo=path_in_repo,
+                repo_id=self.repo_id,
+                repo_type="dataset"
+            )
+            
+            # 一時ファイルを削除
+            try:
+                os.remove(tmp_path)
+            except Exception as e:
+                logger.warning(f"一時ファイル削除エラー: {e}")
+            
+            # URLを構築して返す
+            url = f"https://huggingface.co/{self.repo_id}/resolve/main/{path_in_repo}"
+            logger.info(f"アップロード成功: {url}")
+            return url
+            
+        except Exception as e:
+            logger.error(f"HuggingFace Hubへのアップロード中にエラー: {e}", exc_info=True)
+            return None
+
+# グローバルなアップローダーインスタンスを作成
+hf_uploader = HuggingFaceUploader()
+
 # --- WebDriverプールの実装 ---
 class WebDriverPool:
     """WebDriverインスタンスを再利用するためのプール"""
@@ -126,6 +199,13 @@ class ScreenshotRequest(BaseModel):
     trim_whitespace: bool = True  # 余白トリミングオプション（デフォルト有効）
     style: str = "standard"  # デフォルトはstandard
 
+# --- レスポンスモデル ---
+class ImageUrlResponse(BaseModel):
+    """画像URLのレスポンスモデル"""
+    url: str
+    status: str = "success"
+    message: str = "画像が正常に生成されました"
+
 # HTMLのFont Awesomeレイアウトを改善する関数 - プリロード機能を追加
 def enhance_font_awesome_layout(html_code):
     """Font Awesomeレイアウトを改善し、プリロードタグを追加"""
@@ -602,8 +682,13 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
 
 # --- 並列処理を活用した新しい関数 ---
 def text_to_screenshot_parallel(text: str, extension_percentage: float, temperature: float = 0.3,
-                             trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
-    """テキストをGemini APIでHTMLに変換し、並列処理でスクリーンショットを生成する関数"""
+                             trim_whitespace: bool = True, style: str = "standard") -> tuple:
+    """
+    テキストをGemini APIでHTMLに変換し、並列処理でスクリーンショットを生成する関数
+    
+    Returns:
+        tuple - (PIL.Image, URL) - 生成された画像とHuggingFaceのURL
+    """
     start_time = time.time()
     logger.info("並列処理によるテキスト→スクリーンショット生成を開始")
     
@@ -749,13 +834,17 @@ def text_to_screenshot_parallel(text: str, extension_percentage: float, temperat
                 img = trim_image_whitespace(img, threshold=248, padding=20)
                 logger.info(f"トリミング後のサイズ: {img.width}x{img.height}")
             
+            # 画像をHuggingFaceにアップロード
+            prefix = f"infographic_{style}"
+            image_url = hf_uploader.upload_image(img, prefix=prefix)
+            
             elapsed = time.time() - start_time
-            logger.info(f"並列処理による生成完了。所要時間: {elapsed:.2f}秒")
-            return img
+            logger.info(f"並列処理による生成完了。所要時間: {elapsed:.2f}秒、URL: {image_url}")
+            return img, image_url
             
         except Exception as e:
             logger.error(f"スクリーンショット生成中にエラー: {e}", exc_info=True)
-            return Image.new('RGB', (1, 1), color=(0, 0, 0))
+            return Image.new('RGB', (1, 1), color=(0, 0, 0)), None
         finally:
             # WebDriverプールに戻す
             if driver_from_pool:
@@ -769,15 +858,36 @@ def text_to_screenshot_parallel(text: str, extension_percentage: float, temperat
     
     except Exception as e:
         logger.error(f"並列処理中のエラー: {e}", exc_info=True)
-        return Image.new('RGB', (1, 1), color=(0, 0, 0))  # エラー時は黒画像
+        return Image.new('RGB', (1, 1), color=(0, 0, 0)), None  # エラー時は黒画像とNone URL
 
 # 従来の非並列版も残す（互換性のため）
 def text_to_screenshot(text: str, extension_percentage: float, temperature: float = 0.3,
-                    trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
+                    trim_whitespace: bool = True, style: str = "standard") -> tuple:
     """テキストをGemini APIでHTMLに変換し、スクリーンショットを生成する統合関数（レガシー版）"""
     # 並列処理版を呼び出す
     return text_to_screenshot_parallel(text, extension_percentage, temperature, trim_whitespace, style)
 
+# 新しい関数: HTMLからスクリーンショットを生成し、HuggingFaceにアップロード
+def render_and_upload_screenshot(html_code: str, extension_percentage: float = 10.0, 
+                                trim_whitespace: bool = True, prefix: str = "screenshot") -> tuple:
+    """
+    HTMLコードからスクリーンショットを生成し、HuggingFaceにアップロードする
+    
+    Returns:
+        tuple - (PIL.Image, URL) - 生成された画像とHuggingFaceのURL
+    """
+    try:
+        # スクリーンショット生成
+        img = render_fullpage_screenshot(html_code, extension_percentage, trim_whitespace)
+        
+        # 画像をHuggingFaceにアップロード
+        image_url = hf_uploader.upload_image(img, prefix=prefix)
+        
+        return img, image_url
+    except Exception as e:
+        logger.error(f"スクリーンショット生成とアップロード中にエラー: {e}", exc_info=True)
+        return Image.new('RGB', (1, 1), color=(0, 0, 0)), None
+
 # --- FastAPI Setup ---
 app = FastAPI()
 
@@ -821,59 +931,56 @@ if os.path.exists(cdn_dir):
     app.mount("/cdn", StaticFiles(directory=cdn_dir), name="cdn")
 
 
-# API Endpoint for screenshot generation
+# API Endpoint for screenshot generation - 更新版（URLを返すように変更）
 @app.post("/api/screenshot",
-          response_class=StreamingResponse,
+          response_model=ImageUrlResponse,
           tags=["Screenshot"],
-          summary="Render HTML to Full Page Screenshot",
-          description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, and returns the full-page screenshot as a PNG image.")
+          summary="Render HTML to Full Page Screenshot and Upload to HuggingFace",
+          description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, uploads to HuggingFace, and returns the URL.")
 async def api_render_screenshot(request: ScreenshotRequest):
     """
-    API endpoint to render HTML and return a screenshot.
+    API endpoint to render HTML, upload to HuggingFace, and return the URL.
     """
     try:
         logger.info(f"API request received. Extension: {request.extension_percentage}%")
-        # Run the blocking Selenium code (now using the pooled version)
-        pil_image = render_fullpage_screenshot(
+        
+        # スクリーンショット生成とアップロード
+        pil_image, image_url = render_and_upload_screenshot(
             request.html_code,
             request.extension_percentage,
-            request.trim_whitespace
+            request.trim_whitespace,
+            prefix="screenshot"
         )
 
-        if pil_image.size == (1, 1):
-             logger.error("Screenshot generation failed, returning 1x1 error image.")
-             # Optionally return a proper error response instead of 1x1 image
-             # raise HTTPException(status_code=500, detail="Failed to generate screenshot")
-
-        # Convert PIL Image to PNG bytes
-        img_byte_arr = BytesIO()
-        pil_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)  # Go to the start of the BytesIO buffer
+        if pil_image.size == (1, 1) or not image_url:
+             logger.error("Screenshot generation failed, or upload failed.")
+             raise HTTPException(status_code=500, detail="Failed to generate or upload screenshot")
 
-        logger.info("Returning screenshot as PNG stream.")
-        return StreamingResponse(img_byte_arr, media_type="image/png")
+        # URLを返す
+        logger.info(f"返却URL: {image_url}")
+        return ImageUrlResponse(url=image_url)
 
     except Exception as e:
         logger.error(f"API Error: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
 
-# --- 新しいGemini API連携エンドポイント（並列処理版） ---
+# --- 新しいGemini API連携エンドポイント（並列処理版）- URLを返すよう更新 ---
 @app.post("/api/text-to-screenshot",
-          response_class=StreamingResponse,
+          response_model=ImageUrlResponse,
           tags=["Screenshot", "Gemini"],
-          summary="テキストからインフォグラフィックを生成",
-          description="テキストをGemini APIを使ってHTMLインフォグラフィックに変換し、スクリーンショットとして返します。")
+          summary="テキストからインフォグラフィックを生成しHuggingFaceにアップロード",
+          description="テキストをGemini APIを使ってHTMLインフォグラフィックに変換し、HuggingFaceにアップロードしたURLを返します。")
 async def api_text_to_screenshot(request: GeminiRequest):
     """
-    テキストからHTMLインフォグラフィックを生成してスクリーンショットを返すAPIエンドポイント
+    テキストからHTMLインフォグラフィックを生成してアップロードし、URLを返すAPIエンドポイント
     """
     try:
         logger.info(f"テキスト→スクリーンショットAPIリクエスト受信。テキスト長さ: {len(request.text)}, "
                    f"拡張率: {request.extension_percentage}%, 温度: {request.temperature}, "
                    f"スタイル: {request.style}")
 
-        # 並列処理版を使用
-        pil_image = text_to_screenshot_parallel(
+        # 並列処理版を使用 - 画像とURLを取得
+        pil_image, image_url = text_to_screenshot_parallel(
             request.text,
             request.extension_percentage,
             request.temperature,
@@ -881,16 +988,13 @@ async def api_text_to_screenshot(request: GeminiRequest):
             request.style
         )
 
-        if pil_image.size == (1, 1):
-            logger.error("スクリーンショット生成に失敗しました。1x1エラー画像を返します。")
+        if pil_image.size == (1, 1) or not image_url:
+            logger.error("スクリーンショット生成に失敗したか、アップロードに失敗しました。")
+            raise HTTPException(status_code=500, detail="Failed to generate or upload screenshot")
 
-        # PIL画像をPNGバイトに変換
-        img_byte_arr = BytesIO()
-        pil_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)  # BytesIOバッファの先頭に戻る
-
-        logger.info("スクリーンショットをPNGストリームとして返します。")
-        return StreamingResponse(img_byte_arr, media_type="image/png")
+        # URLを返す
+        logger.info(f"返却URL: {image_url}")
+        return ImageUrlResponse(url=image_url)
 
     except Exception as e:
         logger.error(f"API Error: {e}", exc_info=True)
@@ -901,11 +1005,24 @@ async def api_text_to_screenshot(request: GeminiRequest):
 def process_input(input_mode, input_text, extension_percentage, temperature, trim_whitespace, style):
     """入力モードに応じて適切な処理を行う"""
     if input_mode == "HTML入力":
-        # HTMLモードの場合は��存の処理（スタイルは使わない）
-        return render_fullpage_screenshot(input_text, extension_percentage, trim_whitespace)
+        # HTMLモードの場合はレンダリングとアップロード
+        img, url = render_and_upload_screenshot(
+            input_text, 
+            extension_percentage, 
+            trim_whitespace, 
+            prefix="html_screenshot"
+        )
+        return img, url if url else "アップロード失敗またはURL取得できませんでした"
     else:
-        # テキスト入力モードの場合はGemini APIを使用（並列処理版）
-        return text_to_screenshot_parallel(input_text, extension_percentage, temperature, trim_whitespace, style)
+        # テキスト入力モードの場合はGemini API使用（並列処理版）
+        img, url = text_to_screenshot_parallel(
+            input_text, 
+            extension_percentage, 
+            temperature, 
+            trim_whitespace, 
+            style
+        )
+        return img, url if url else "アップロード失敗またはURL取得できませんでした"
 
 # Gradio UIの定義
 with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Base()) as iface:
@@ -965,7 +1082,13 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
     )
 
     submit_btn = gr.Button("生成")
-    output_image = gr.Image(type="pil", label="ページ全体のスクリーンショット")
+    
+    # 出力部分をRowで分ける
+    with gr.Row():
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="pil", label="ページ全体のスクリーンショット")
+        with gr.Column(scale=1):
+            output_url = gr.Textbox(label="画像URL（HuggingFace）", info="生成された画像のURLです。このURLを使用して画像にアクセスできます。")
 
     # 入力モード変更時のイベント処理（テキストモード時のみ温度スライダーとスタイルドロップダウンを表示）
     def update_controls_visibility(mode):
@@ -986,20 +1109,22 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
     submit_btn.click(
         fn=process_input,
         inputs=[input_mode, input_text, extension_percentage, temperature, trim_whitespace, style_dropdown],
-        outputs=output_image
+        outputs=[output_image, output_url]
     )
 
     # 環境変数情報を表示
     gemini_model = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
+    hf_repo = os.environ.get("HF_REPO_ID", "tomo2chin2/SUPER_TENSAI_JIN")
     gr.Markdown(f"""
     ## APIエンドポイント
-    - `/api/screenshot` - HTMLコードからスクリーンショットを生成
-    - `/api/text-to-screenshot` - テキストからインフォグラフィックスクリーンショットを生成
+    - `/api/screenshot` - HTMLコードからスクリーンショットを生成し、URLを返します
+    - `/api/text-to-screenshot` - テキストからインフォグラフィックスクリーンショットを生成し、URLを返します
 
     ## 設定情報
     - 使用モデル: {gemini_model} (環境変数 GEMINI_MODEL で変更可能)
-    - 対応スタイル: standard, cute, resort, cool, dental
+    - HuggingFaceリポジトリ: {hf_repo} (環境変数 HF_REPO_ID で変更可能)
     - WebDriverプール最大数: {driver_pool.max_drivers} (環境変数 MAX_WEBDRIVERS で変更可能)
+    - 対応スタイル: standard, cute, resort, cool, dental
     """)
 
 # --- Mount Gradio App onto FastAPI ---