Update app.py

app.py

@@ -1,33 +1,11 @@
 import gradio as gr
 import pandas as pd
 import os
-import json
-import base64
 import requests
 from io import BytesIO
 from PIL import Image
 from urllib.parse import urlparse, parse_qs
-from google.cloud import aiplatform
-from google.api_core.client_options import ClientOptions
-from google.oauth2 import service_account
-
-# 由於您的環境可能沒有服務帳號金鑰，我們將直接使用 GEMINI_API_KEY
-# 如果在 Hugging Face Space 中，這通常會從環境變數中自動加載
-
-# 初始化 AI Platform 服務客戶端
-def initialize_ai_platform(api_key):
-    # 請替換為您的 GCP 專案 ID 和區域
-    PROJECT_ID = "YOUR_GCP_PROJECT_ID"  # 這裡需要您提供真實的GCP專案ID
-    LOCATION = "us-central1"
-
-    # 這裡的認證方式取決於您的環境。
-    # 創建一個不帶服務帳號憑證的客戶端，這會讓它嘗試從環境中自動尋找憑證
-    try:
-        aiplatform.init(project=PROJECT_ID, location=LOCATION)
-        return True
-    except Exception as e:
-        print(f"Error initializing AI Platform: {e}")
-        return False
+import google.generativeai as genai
 
 # 全局變數來儲存 API 金鑰
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
@@ -43,144 +21,106 @@ print(f"Debug: Top-level Loaded GEMINI_API_KEY (first 5 chars): {GEMINI_API_KEY[
 if not GEMINI_API_KEY:
     raise ValueError("ERROR: GEMINI_API_KEY environment variable is not set. Please set it correctly.")
 
-# 這裡我們將創建一個通用客戶端來處理 API 呼叫
-# 由於 'google-cloud-aiplatform' 函式庫主要設計用於 GCP 環境，我們將使用 `requests`
-# 來模擬直接呼叫 API，以適應您在 Hugging Face Space 的需求。
-class NanoBananaImageGenerator:
-    def __init__(self, api_key):
-        self.api_key = api_key
-        self.endpoint_url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent"
-        
-    def _load_image_from_url(self, url):
-        # 圖片下載邏輯與之前相同，以確保兼容性
-        try:
-            headers = {'User-Agent': 'Mozilla/5.0'}
-            response = requests.get(url, timeout=20, headers=headers)
-            response.raise_for_status()
-            image = Image.open(BytesIO(response.content)).convert("RGB")
-            return image
-        except requests.exceptions.HTTPError as e:
-            print(f"Error downloading image from {url}: HTTP Error {e.response.status_code}")
-            return None
-        except Exception as e:
-            print(f"An unexpected error occurred: {e}")
-            return None
+# 配置 Gemini API
+genai.configure(api_key=GEMINI_API_KEY)
 
-    def _image_to_base64(self, pil_image):
-        img_byte_arr = BytesIO()
-        pil_image.save(img_byte_arr, format='PNG')
-        img_bytes = img_byte_arr.getvalue()
-        return base64.b64encode(img_bytes).decode('utf-8')
+def load_image_from_url(url: str):
+    """從 URL 下載圖片並以 PIL Image 格式回傳。"""
+    try:
+        headers = {'User-Agent': 'Mozilla/5.0'}
+        response = requests.get(url, timeout=20, headers=headers)
+        response.raise_for_status()  # 檢查請求是否成功
+        image = Image.open(BytesIO(response.content)).convert("RGB")
+        print(f"Debug: Successfully loaded image from URL: {url}")
+        return image
+    except requests.exceptions.HTTPError as e:
+        print(f"Error downloading image from {url}: HTTP Error {e.response.status_code}")
+        return None
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        return None
+
+def build_prompt_for_operation(prompt, has_references=False, aspect_ratio="1:1"):
+    """根據輸入構建完整的提示詞。"""
+    aspect_instructions = {
+        "1:1": "square format",
+        "16:9": "widescreen landscape format",
+        "9:16": "portrait format",
+        "4:3": "standard landscape format",
+        "3:4": "standard portrait format"
+    }
+    base_quality = "Generate a high-quality, photorealistic image"
+    format_instruction = f"in {aspect_instructions.get(aspect_ratio, 'square format')}"
+    
+    final_prompt = f"{base_quality} inspired by the style and elements of the reference images. {prompt}. {format_instruction}."
+    if not has_references:
+        final_prompt = f"{base_quality} of: {prompt}. {format_instruction}."
+        
+    return final_prompt
 
-    def build_prompt_for_operation(self, prompt, has_references=False, aspect_ratio="1:1"):
-        # 提示詞生成邏輯與之前相同
-        aspect_instructions = {
-            "1:1": "square format",
-            "16:9": "widescreen landscape format",
-            "9:16": "portrait format",
-            "4:3": "standard landscape format",
-            "3:4": "standard portrait format"
-        }
-        base_quality = "Generate a high-quality, photorealistic image"
-        format_instruction = f"in {aspect_instructions.get(aspect_ratio, 'square format')}"
+def call_gemini_api(prompt, images):
+    """使用官方函式庫呼叫 Gemini API。"""
+    operation_log = ""
+    try:
+        # 使用官方 genai.generate_content 函式
+        response = genai.generate_content(
+            contents=[prompt] + images, # 圖片列表直接作為輸入
+            model="gemini-1.5-pro-latest" # 使用官方推薦的多模態模型
+        )
         
-        final_prompt = f"{base_quality} inspired by the style and elements of the reference images. {prompt}. {format_instruction}."
-        if not has_references:
-            final_prompt = f"{base_quality} of: {prompt}. {format_instruction}."
+        print(f"Debug: Full API Response: {response.text}")
         
-        return final_prompt
+        # 檢查是否有安全政策問題
+        if 'prompt_feedback' in response:
+            if 'safety_ratings' in response['prompt_feedback']:
+                for rating in response['prompt_feedback']['safety_ratings']:
+                    if rating['block_reason'] != 'NONE':
+                        operation_log += f"API 被安全政策阻止。原因：{rating['block_reason']}\n"
+                        return None, operation_log
         
-    def call_nano_banana_api(self, prompt, images, temperature=0.7):
-        operation_log = ""
-        try:
-            # 構建請求主體
-            contents = [
-                {"text": prompt}
-            ]
-            for img in images:
-                contents.append({
-                    "inline_data": {
-                        "mime_type": "image/png",
-                        "data": self._image_to_base64(img)
-                    }
-                })
-
-            payload = {
-                "contents": {"parts": contents},
-                "generation_config": {
-                    "temperature": temperature,
-                    "response_modalities": ["Text", "Image"]
-                }
-            }
-            
-            headers = {
-                "Content-Type": "application/json",
-                "x-goog-api-key": self.api_key
-            }
-
-            print(f"Debug: Attempting to call Gemini API via HTTP request.")
-            response = requests.post(self.endpoint_url, headers=headers, json=payload, timeout=60)
-            response.raise_for_status()
-            
-            api_response = response.json()
-            print(f"Debug: Full API Response: {json.dumps(api_response, indent=2)}")
-
-            # 解析回應
-            candidates = api_response.get('candidates', [])
-            if not candidates:
-                safety_ratings = api_response.get('prompt_feedback', {}).get('safety_ratings', [])
-                if safety_ratings:
-                    operation_log += f"API 被安全政策阻止。原因：{safety_ratings}\n"
-                else:
-                    operation_log += "API 回應中未找到候選者，可能的原因是內部錯誤或無效請求。\n"
-                return [], operation_log
-            
-            image_parts = []
-            for candidate in candidates:
-                if 'content' in candidate and 'parts' in candidate['content']:
-                    for part in candidate['content']['parts']:
-                        if 'inline_data' in part and 'data' in part['inline_data']:
-                            image_parts.append(base64.b64decode(part['inline_data']['data']))
-                        elif 'text' in part:
-                            operation_log += f"API 回應文字：{part['text']}\n"
-                
-            if not image_parts:
-                operation_log += "API 回應沒有包含任何圖像數據，請檢查輸入內容是否違反安全政策。\n"
-            else:
-                operation_log += f"成功生成 {len(image_parts)} 張圖片。\n"
-            
+        # 處理回應
+        if response.text is not None and "data:image" in response.text:
+            # 這是內嵌的圖片 Base64 字串，需要解碼
+            base64_string = response.text.split(',')[1]
+            image_data = base64.b64decode(base64_string)
+            image_parts = [image_data]
+            operation_log += f"成功生成 {len(image_parts)} 張圖片。\n"
             return image_parts, operation_log
-            
-        except requests.exceptions.HTTPError as e:
-            operation_log = f"API 呼叫失敗，HTTP Error: {e.response.status_code} - {e.response.text}\n"
-            return [], operation_log
-        except Exception as e:
-            operation_log = f"意外錯誤: {type(e).__name__} - {str(e)}\n"
-            return [], operation_log
+        elif response.text is not None:
+            # 如果回傳的是文字
+            operation_log += f"API 回應文字：{response.text}\n"
+            return None, operation_log
+        else:
+            # 沒有任何回傳
+            operation_log += "API 回應沒有包含任何圖像或文字數據。\n"
+            return None, operation_log
+
+    except Exception as e:
+        operation_log = f"意外錯誤: {type(e).__name__} - {str(e)}\n"
+        return None, operation_log
 
 def generate_image(white_background_url, reference_image_url, prompt):
-    image_generator = NanoBananaImageGenerator(api_key=GEMINI_API_KEY)
-    
+    """Gradio 介面呼叫的主函式。"""
     if not GEMINI_API_KEY:
         return None, "Error: GEMINI_API_KEY is not set."
 
-    wb_image = image_generator._load_image_from_url(white_background_url)
-    ref_image = image_generator._load_image_from_url(reference_image_url)
+    wb_image = load_image_from_url(white_background_url)
+    ref_image = load_image_from_url(reference_image_url)
 
-    if not wb_image:
-        return None, "Error: Failed to load white background image from URL."
-    if not ref_image:
-        return None, "Error: Failed to load reference image from URL."
+    if wb_image is None or ref_image is None:
+        return None, "Error: One or more images failed to load from URL."
 
     images = [wb_image, ref_image]
-    final_prompt = image_generator.build_prompt_for_operation(prompt, has_references=True)
+    final_prompt = build_prompt_for_operation(prompt, has_references=True)
 
-    generated_images_binary, operation_log = image_generator.call_nano_banana_api(final_prompt, images)
+    generated_images_binary, operation_log = call_gemini_api(final_prompt, images)
 
     if generated_images_binary:
         output_dir = "generated_images"
         os.makedirs(output_dir, exist_ok=True)
-        output_path = os.path.join(output_dir, f"generated_{len(os.listdir(output_dir)) + 1}.png")
+        # 使用時間戳或唯一ID來確保檔名唯一
+        import time
+        output_path = os.path.join(output_dir, f"generated_{int(time.time())}.png")
         with open(output_path, "wb") as f:
             f.write(generated_images_binary[0])
         return output_path, operation_log
@@ -188,6 +128,7 @@ def generate_image(white_background_url, reference_image_url, prompt):
         return None, operation_log
 
 def read_google_sheet(sheet_url):
+    """從 Google Sheet 讀取資料。"""
     try:
         def build_csv_url(url: str) -> str:
             parsed = urlparse(url)
@@ -214,6 +155,7 @@ def read_google_sheet(sheet_url):
         raise gr.Error(f"Error reading Google Sheet: {e}")
 
 def process_sheet_data(sheet_url):
+    """處理試算表資料，為 Gradio DataFrame 準備。"""
     try:
         df = read_google_sheet(sheet_url)
 
@@ -221,33 +163,35 @@ def process_sheet_data(sheet_url):
             error_msg = f"Error: Google Sheet has only {df.shape[1]} columns, but 3 are expected (White Background URL, Reference Image URL, Prompt)."
             raise gr.Error(error_msg)
 
-        white_background_urls = df.iloc[:, 0].tolist()
-        reference_image_urls = df.iloc[:, 1].tolist()
-        prompts = df.iloc[:, 2].tolist()
-
-        data = []
-        for i, (wb, ref, p) in enumerate(zip(white_background_urls, reference_image_urls, prompts)):
-            if pd.notna(wb) and pd.notna(ref) and pd.notna(p):
-                data.append([i, wb, ref, p])
+        # 這裡使用 to_dict('records') 來處理，以確保後續 Pandas 處理時的兼容性
+        data = df.to_dict('records')
+        data_list = []
+        for i, row in enumerate(data):
+            if pd.notna(row.iloc[0]) and pd.notna(row.iloc[1]) and pd.notna(row.iloc[2]):
+                data_list.append([i, row.iloc[0], row.iloc[1], row.iloc[2]])
 
-        return data
+        return data_list
     except Exception as e:
         raise gr.Error(f"Error processing sheet data: {e}")
 
 def generate_image_for_row(row_index, dataframe_data):
-    if not (0 <= row_index < len(dataframe_data)):
+    """根據 Gradio DataFrame 的行索引生成圖片。"""
+    # 將 Gradio 的 dataframe_data 轉換回 pandas DataFrame
+    df = pd.DataFrame(dataframe_data, columns=["Index", "白背圖URL", "參考圖URL", "提示詞"])
+    
+    if not (0 <= row_index < len(df)):
         return None, "Error: Invalid row index."
 
-    row = dataframe_data.iloc[row_index]
-    white_background_url = row.iloc[1]
-    reference_image_url = row.iloc[2]
-    prompt = row.iloc[3]
+    row = df.iloc[int(row_index)]
+    white_background_url = row['白背圖URL']
+    reference_image_url = row['參考圖URL']
+    prompt = row['提示詞']
 
     return generate_image(white_background_url, reference_image_url, prompt)
 
 if __name__ == "__main__":
     with gr.Blocks() as demo:
-        gr.Markdown("# AutoLS Gradio Image Generator (Powered by AI Platform)")
+        gr.Markdown("# AutoLS Gradio Image Generator (Powered by Gemini API)")
         gr.Markdown("輸入 Google Sheet 網址來處理圖像生成請求。")
 
         sheet_url_input = gr.Textbox(label="Google Sheet URL", value="https://docs.google.com/spreadsheets/d/1G3olHxydDIbnyXdh5nnw5TG0akZFeMeYm-25JmCGDLg/edit?gid=0#gid=0")
@@ -273,14 +217,15 @@ if __name__ == "__main__":
             inputs=sheet_url_input,
             outputs=output_dataframe
         ).success(
-            fn=lambda x: pd.DataFrame(x, columns=["Index", "白背圖URL", "參考圖URL", "提示詞"]),
+            fn=lambda x: x,
             inputs=output_dataframe,
             outputs=processed_df_state
         )
-
+        
+        # 修正 click 觸發器，確保傳入正確的參數
         generate_selected_button.click(
             fn=generate_image_for_row,
-            inputs=[row_index_input, processed_df_state],
+            inputs=[row_index_input, output_dataframe],
             outputs=[generated_image_output, operation_log_output]
         )