Update app.py

app.py

@@ -12,6 +12,44 @@ import copy
 import random
 import time
 from transformers import pipeline
+import sqlite3
+from datetime import datetime
+
+# 데이터베이스 초기화
+def init_db():
+    conn = sqlite3.connect('gallery.db')
+    c = conn.cursor()
+    c.execute('''CREATE TABLE IF NOT EXISTS images
+                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
+                  model_name TEXT,
+                  prompt TEXT,
+                  image_path TEXT,
+                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
+    conn.commit()
+    conn.close()
+
+# 이미지 저장 함수
+def save_image(image, model_name, prompt):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"gallery_{timestamp}.png"
+    image.save(os.path.join("gallery", filename))
+    
+    conn = sqlite3.connect('gallery.db')
+    c = conn.cursor()
+    c.execute("INSERT INTO images (model_name, prompt, image_path) VALUES (?, ?, ?)",
+              (model_name, prompt, filename))
+    conn.commit()
+    conn.close()
+
+# 갤러리 이미지 로드 함수
+def load_gallery_images():
+    conn = sqlite3.connect('gallery.db')
+    c = conn.cursor()
+    c.execute("SELECT model_name, prompt, image_path FROM images ORDER BY created_at DESC")
+    rows = c.fetchall()
+    conn.close()
+    
+    return [(os.path.join("gallery", row[2]), f"{row[0]}: {row[1]}") for row in rows]
 
 # 번역 모델 초기화
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
@@ -143,7 +181,10 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         final_image = image
         progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
         yield image, seed, gr.update(value=progress_bar, visible=True), original_prompt, english_prompt
-        
+    
+    # 갤러리에 이미지 저장
+    save_image(final_image, selected_lora['title'], original_prompt)
+    
     yield final_image, seed, gr.update(value=progress_bar, visible=False), original_prompt, english_prompt
 
 
@@ -222,6 +263,7 @@ def add_custom_lora(custom_lora):
 def remove_custom_lora():
     return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
 
+
 run_lora.zerogpu = True
 
 css = """
@@ -230,52 +272,70 @@ footer {
 }
 """
 
+# 갤러리 디렉토리 생성
+if not os.path.exists('gallery'):
+    os.makedirs('gallery')
+
+# 데이터베이스 초기화
+init_db()
 
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
 
     selected_index = gr.State(None)
-    with gr.Row():
-        with gr.Column(scale=3):
-            prompt = gr.Textbox(label="프롬프트", lines=1, placeholder="LoRA를 선택한 후 프롬프트를 입력하세요 (한글 또는 영어)")
-        with gr.Column(scale=1, elem_id="gen_column"):
-            generate_button = gr.Button("생성", variant="primary", elem_id="gen_btn")
-    with gr.Row():
-        with gr.Column():
-            selected_info = gr.Markdown("")
-            gallery = gr.Gallery(
-                [(item["image"], item["title"]) for item in loras],
-                label="LoRA 갤러리",
-                allow_preview=False,
+    
+    with gr.Tabs():
+        with gr.TabItem("생성"):
+            with gr.Row():
+                with gr.Column(scale=3):
+                    prompt = gr.Textbox(label="프롬프트", lines=1, placeholder="LoRA를 선택한 후 프롬프트를 입력하세요 (한글 또는 영어)")
+                with gr.Column(scale=1, elem_id="gen_column"):
+                    generate_button = gr.Button("생성", variant="primary", elem_id="gen_btn")
+            with gr.Row():
+                with gr.Column():
+                    selected_info = gr.Markdown("")
+                    gallery = gr.Gallery(
+                        [(item["image"], item["title"]) for item in loras],
+                        label="LoRA 갤러리",
+                        allow_preview=False,
+                        columns=3,
+                        elem_id="gallery"
+                    )
+                    with gr.Group():
+                        custom_lora = gr.Textbox(label="커스텀 LoRA", info="LoRA Hugging Face 경로", placeholder="multimodalart/vintage-ads-flux")
+                        gr.Markdown("[FLUX LoRA 목록 확인](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
+                    custom_lora_info = gr.HTML(visible=False)
+                    custom_lora_button = gr.Button("커스텀 LoRA 제거", visible=False)
+                with gr.Column():
+                    progress_bar = gr.Markdown(elem_id="progress",visible=False)
+                    result = gr.Image(label="생성된 이미지")
+                    original_prompt_display = gr.Textbox(label="원본 프롬프트")
+                    english_prompt_display = gr.Textbox(label="영어 프롬프트")
+
+            with gr.Row():
+                with gr.Accordion("고급 설정", open=False):
+                    with gr.Column():
+                        with gr.Row():
+                            cfg_scale = gr.Slider(label="CFG 스케일", minimum=1, maximum=20, step=0.5, value=3.5)
+                            steps = gr.Slider(label="스텝", minimum=1, maximum=50, step=1, value=28)
+                        
+                        with gr.Row():
+                            width = gr.Slider(label="너비", minimum=256, maximum=1536, step=64, value=1024)
+                            height = gr.Slider(label="높이", minimum=256, maximum=1536, step=64, value=1024)
+                        
+                        with gr.Row():
+                            randomize_seed = gr.Checkbox(True, label="시드 무작위화")
+                            seed = gr.Slider(label="시드", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
+                            lora_scale = gr.Slider(label="LoRA 스케일", minimum=0, maximum=3, step=0.01, value=0.95)
+
+        with gr.TabItem("갤러리"):
+            gallery_images = gr.Gallery(
+                load_gallery_images(),
+                label="생성된 이미지 갤러리",
                 columns=3,
-                elem_id="gallery"
+                rows=3,
+                height="auto"
             )
-            with gr.Group():
-                custom_lora = gr.Textbox(label="커스텀 LoRA", info="LoRA Hugging Face 경로", placeholder="multimodalart/vintage-ads-flux")
-                gr.Markdown("[FLUX LoRA 목록 확인](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
-            custom_lora_info = gr.HTML(visible=False)
-            custom_lora_button = gr.Button("커스텀 LoRA 제거", visible=False)
-        with gr.Column():
-            progress_bar = gr.Markdown(elem_id="progress",visible=False)
-            result = gr.Image(label="생성된 이미지")
-            original_prompt_display = gr.Textbox(label="원본 프롬프트")
-            english_prompt_display = gr.Textbox(label="영어 프롬프트")
-
-    with gr.Row():
-        with gr.Accordion("고급 설정", open=False):
-            with gr.Column():
-                with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG 스케일", minimum=1, maximum=20, step=0.5, value=3.5)
-                    steps = gr.Slider(label="스텝", minimum=1, maximum=50, step=1, value=28)
-                
-                with gr.Row():
-                    width = gr.Slider(label="너비", minimum=256, maximum=1536, step=64, value=1024)
-                    height = gr.Slider(label="높이", minimum=256, maximum=1536, step=64, value=1024)
-                
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(True, label="시드 무작위화")
-                    seed = gr.Slider(label="시드", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-                    lora_scale = gr.Slider(label="LoRA 스케일", minimum=0, maximum=3, step=0.01, value=0.95)
-
+            refresh_button = gr.Button("갤러리 새로고침")
 
     gallery.select(
         update_selection,
@@ -299,5 +359,10 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
         outputs=[result, seed, progress_bar, original_prompt_display, english_prompt_display]
     )
 
+    refresh_button.click(
+        lambda: gr.update(value=load_gallery_images()),
+        outputs=[gallery_images]
+    )
+
 app.queue()
 app.launch()