Update app.py

app.py

@@ -1,14 +1,21 @@
-import gradio as gr
-from gradio_client import Client
 import os
+from dotenv import load_dotenv
 import logging
+from notion_client import Client
+import gradio as gr
+from gradio_client import Client as GradioClient
 import requests
-import json
 from datetime import datetime
 
 # 로깅 설정
 logging.basicConfig(level=logging.INFO)
 
+# 환경 변수 로드
+load_dotenv()
+
+# 로깅 설정
+logging.basicConfig(level=logging.INFO)
+
 # API 클라이언트 설정
 api_client = Client("http://211.233.58.202:7960/")
 
@@ -17,64 +24,37 @@ NOTION_API_KEY = "secret_MpVfJphbfo4599fdczYfMYKNOpyzCcvkhhzk3lgTfVk"
 NOTION_DATABASE_ID = "88c9bdad-cb20-4412-9af7-7d5932e1a82a"
 NOTION_API_URL = "https://api.notion.com/v1"
 
-headers = {
-    "Authorization": f"Bearer {NOTION_API_KEY}",
-    "Content-Type": "application/json",
-    "Notion-Version": "2022-06-28"  # 최신 버전으로 업데이트
-}
+# Notion 클라이언트 초기화
+notion = Client(auth=NOTION_API_KEY)
 
-def upload_to_notion(prompt, image_url):
-    current_time = datetime.now().isoformat()
-    
-    data = {
-        "parent": {"database_id": NOTION_DATABASE_ID},
-        "properties": {
-            "Name": {
-                "title": [
-                    {
-                        "text": {
-                            "content": f"Generated Image - {current_time}"
-                        }
-                    }
-                ]
-            },
-            "Prompt": {
-                "rich_text": [
-                    {
-                        "text": {
-                            "content": prompt
-                        }
-                    }
-                ]
-            }
-        },
-        "children": [
-            {
-                "object": "block",
-                "type": "image",
-                "image": {
-                    "type": "external",
-                    "external": {
-                        "url": image_url
-                    }
-                }
-            }
-        ]
-    }
+# Gradio API 클라이언트 설정
+gradio_client = GradioClient("http://211.233.58.202:7960/")
 
-    try:
-        response = requests.post(f"{NOTION_API_URL}/pages", headers=headers, json=data)
-        response.raise_for_status()
-        logging.info("Successfully added to Notion")
-        return response.json()["url"]
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Failed to add to Notion: {e}")
-        logging.error(f"Response content: {response.text}")
-        logging.error(f"Request payload: {json.dumps(data, indent=2)}")
-        logging.error(f"Headers: {headers}")  # API 키는 로그에 남기지 않도록 주의하세요
-        raise Exception(f"Failed to add to Notion: {e}")
+def create_notion_database():
+    new_database = notion.databases.create(
+        parent={"type": "page_id", "page_id": NOTION_PAGE_ID},
+        title=[{"type": "text", "text": {"content": "Generated Images"}}],
+        properties={
+            "Name": {"title": {}},
+            "Prompt": {"rich_text": {}},
+            "Image": {"url": {}},
+            "Created At": {"date": {}}
+        }
+    )
+    return new_database.id
+
+def add_page_to_notion_database(database_id, name, prompt, image_url):
+    new_page = notion.pages.create(
+        parent={"database_id": database_id},
+        properties={
+            "Name": {"title": [{"text": {"content": name}}]},
+            "Prompt": {"rich_text": [{"text": {"content": prompt}}]},
+            "Image": {"url": image_url},
+            "Created At": {"date": {"start": datetime.now().isoformat()}}
+        }
+    )
+    return new_page.url
 
-        
 def respond(message, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     logging.info(f"Received message: {message}, seed: {seed}, randomize_seed: {randomize_seed}, "
                  f"width: {width}, height: {height}, guidance_scale: {guidance_scale}, "
@@ -82,7 +62,7 @@ def respond(message, seed, randomize_seed, width, height, guidance_scale, num_in
     
     try:
         # 이미지 생성 요청
-        result = api_client.predict(
+        result = gradio_client.predict(
             prompt=message,
             seed=seed,
             randomize_seed=randomize_seed,
@@ -97,19 +77,39 @@ def respond(message, seed, randomize_seed, width, height, guidance_scale, num_in
         # 결과 확인 및 처리
         if isinstance(result, tuple) and len(result) >= 1:
             image_path = result[0]
-            # 이미지 URL 생성 (실제 서버 URL로 변경 필요)
             image_url = f"http://211.233.58.202:7960/file={image_path}"
             
-            # Notion에 업로드
-            notion_url = upload_to_notion(message, image_url)
+            # Notion 데이터베이스 생성 (이미 존재한다면 기존 ID 사용)
+            database_id = os.getenv("NOTION_DATABASE_ID")
+            if not database_id:
+                database_id = create_notion_database()
+                os.environ["NOTION_DATABASE_ID"] = database_id
+            
+            # Notion에 페이지 추가
+            page_url = add_page_to_notion_database(
+                database_id,
+                f"Generated Image - {datetime.now().isoformat()}",
+                message,
+                image_url
+            )
             
-            return image_url, notion_url
+            logging.info(f"Added to Notion: {page_url}")
+            return image_url, page_url
         else:
             raise ValueError("Unexpected API response format")
     except Exception as e:
         logging.error("Error during API request: %s", str(e))
         return "Failed to generate image due to an error.", None
 
+
+
+
+
+
+
+
+
+
 def use_prompt(prompt):
     return prompt
 
@@ -143,11 +143,14 @@ examples = [
     ["A fantasy map of a fictional world, with detailed terrain and cities.", "q19.webp"]
 ]
 
+
+
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
+
+    
     with gr.Row():
         input_text = gr.Textbox(label="Enter your prompt for image generation")
         output_image = gr.Image(label="Generated Image")
-        notion_link = gr.Textbox(label="Notion Page Link")
     
     with gr.Row():
         seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="Seed", value=123)
@@ -179,8 +182,8 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
     input_text.submit(
         fn=respond, 
         inputs=[input_text, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], 
-        outputs=[output_image, notion_link]
+        outputs=output_image
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()