Update app.py

app.py

@@ -1,61 +1,57 @@
+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/")
 
-# Notion API 설정
-NOTION_API_KEY = "secret_MpVfJphbfo4599fdczYfMYKNOpyzCcvkhhzk3lgTfVk"
-NOTION_DATABASE_ID = "88c9bdad-cb20-4412-9af7-7d5932e1a82a"
-NOTION_API_URL = "https://api.notion.com/v1"
-
-# Notion 클라이언트 초기화
-notion = Client(auth=NOTION_API_KEY)
-
-# Gradio API 클라이언트 설정
-gradio_client = GradioClient("http://211.233.58.202:7960/")
-
-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
+# 갤러리 저장 디렉토리 설정
+GALLERY_DIR = "gallery"
+GALLERY_JSON = "gallery.json"
+
+# 갤러리 디렉토리 생성
+os.makedirs(GALLERY_DIR, exist_ok=True)
 
+def save_to_gallery(image_path, prompt):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    new_image_path = os.path.join(GALLERY_DIR, f"{timestamp}.png")
+    
+    # 이미지 파일 복사
+    with open(image_path, "rb") as src, open(new_image_path, "wb") as dst:
+        dst.write(src.read())
+    
+    # 갤러리 정보 저장
+    gallery_info = {
+        "image": new_image_path,
+        "prompt": prompt,
+        "timestamp": timestamp
+    }
+    
+    if os.path.exists(GALLERY_JSON):
+        with open(GALLERY_JSON, "r") as f:
+            gallery = json.load(f)
+    else:
+        gallery = []
+    
+    gallery.append(gallery_info)
+    
+    with open(GALLERY_JSON, "w") as f:
+        json.dump(gallery, f, indent=2)
+    
+    return new_image_path
+
+def load_gallery():
+    if os.path.exists(GALLERY_JSON):
+        with open(GALLERY_JSON, "r") as f:
+            gallery = json.load(f)
+        return [(item["image"], item["prompt"]) for item in reversed(gallery)]
+    return []
 
 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}, "
@@ -64,7 +60,7 @@ def respond(message, seed, randomize_seed, width, height, guidance_scale, num_in
     
     try:
         # 이미지 생성 요청
-        result = gradio_client.predict(
+        result = api_client.predict(
             prompt=message,
             seed=seed,
             randomize_seed=randomize_seed,
@@ -79,33 +75,13 @@ def respond(message, seed, randomize_seed, width, height, guidance_scale, num_in
         # 결과 확인 및 처리
         if isinstance(result, tuple) and len(result) >= 1:
             image_path = result[0]
-            image_url = f"http://211.233.58.202:7960/file={image_path}"
-            
-            # Notion 데이터베이스 생성 (이미 존재한다면 기존 ID 사용)
-            database_id = NOTION_DATABASE_ID  # 이 변수를 미리 정의해야 합니다
-            if not database_id:
-                database_id = create_notion_database()
-                NOTION_DATABASE_ID = database_id  # 전역 변수로 저장
-            
-            # Notion에 페이지 추가
-            page_url = add_page_to_notion_database(
-                database_id,
-                f"Generated Image - {datetime.now().isoformat()}",
-                message,
-                image_url
-            )
-            
-            logging.info(f"Added to Notion: {page_url}")
-            return image_url, page_url
+            saved_image_path = save_to_gallery(image_path, message)
+            return saved_image_path
         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
+        return "Failed to generate image due to an error."
 
 css = """
 footer {
@@ -113,6 +89,8 @@ footer {
 }
 """
 
+
+
 # 이미지 생성을 위한 예제 프롬프트
 examples = [
     ["A glamorous young woman with long, wavy blonde hair and smokey eye makeup, posing in a luxury hotel room. She’s wearing a sparkly gold cocktail dress and holding up a white card with 'openfree.ai' written on it in elegant calligraphy. Soft, warm lighting creates a luxurious atmosphere. ", "q1.webp"],
@@ -137,46 +115,68 @@ examples = [
     ["A fantasy map of a fictional world, with detailed terrain and cities.", "q19.webp"]
 ]
 
-
+def use_prompt(prompt):
+    return prompt
 
 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")
-    
-    with gr.Row():
-        seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="Seed", value=123)
-        randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
-    
-    with gr.Row():
-        width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=1024)
-        height = gr.Slider(minimum=256, maximum=1024, step=64, label="Height", value=576)
-    
-    with gr.Row():
-        guidance_scale = gr.Slider(minimum=1, maximum=20, step=0.1, label="Guidance Scale", value=5)
-        num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
-    
-    with gr.Row():
-        for prompt, image_file in examples:
-            with gr.Column():
-                gr.Image(image_file, label=prompt[:50] + "...")  # 프롬프트의 처음 50자만 표시
-                gr.Button("Use this prompt").click(
-                    fn=use_prompt,
-                    inputs=[],
-                    outputs=input_text,
-                    api_name=False
-                ).then(
-                    lambda x=prompt: x,
-                    inputs=[],
-                    outputs=input_text
-                )
+    with gr.Tab("Generate"):
+        with gr.Row():
+            input_text = gr.Textbox(label="Enter your prompt for image generation")
+            output_image = gr.Image(label="Generated Image")
+        
+        with gr.Row():
+            seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="Seed", value=123)
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+        
+        with gr.Row():
+            width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=1024)
+            height = gr.Slider(minimum=256, maximum=1024, step=64, label="Height", value=576)
+        
+        with gr.Row():
+            guidance_scale = gr.Slider(minimum=1, maximum=20, step=0.1, label="Guidance Scale", value=5)
+            num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
+        
+        with gr.Row():
+            for prompt, image_file in examples:
+                with gr.Column():
+                    gr.Image(image_file, label=prompt[:50] + "...")  # 프롬프트의 처음 50자만 표시
+                    gr.Button("Use this prompt").click(
+                        fn=use_prompt,
+                        inputs=[],
+                        outputs=input_text,
+                        api_name=False
+                    ).then(
+                        lambda x=prompt: x,
+                        inputs=[],
+                        outputs=input_text
+                    )
     
+    with gr.Tab("Gallery"):
+        gallery = gr.Gallery(
+            label="Generated Images",
+            show_label=False,
+            elem_id="gallery",
+            columns=[5],
+            rows=[3],
+            object_fit="contain",
+            height="auto"
+        )
+        refresh_btn = gr.Button("Refresh Gallery")
+
+    def update_gallery():
+        return load_gallery()
+
+    refresh_btn.click(fn=update_gallery, inputs=None, outputs=gallery)
+    demo.load(fn=update_gallery, inputs=None, outputs=gallery)
+
     input_text.submit(
         fn=respond, 
         inputs=[input_text, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], 
         outputs=output_image
+    ).then(
+        fn=update_gallery,
+        inputs=None,
+        outputs=gallery
     )
 
 if __name__ == "__main__":