Update app.py

app.py

@@ -26,6 +26,8 @@ from PIL import Image
 from gradio_client import Client, handle_file
 import uuid
 
+import random
+from datetime import datetime
 
 def clear_memory():
     """메모리 정리 함수"""
@@ -702,8 +704,64 @@ def update_controls(bg_prompt):
         gr.update(visible=is_visible),  # object_controls
     ]
 
+
+# 저장 디렉토리 설정
+SAVE_DIR = "saved_images"
+if not os.path.exists(SAVE_DIR):
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def save_generated_image(image, prompt):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"{timestamp}_{unique_id}.png"
+    filepath = os.path.join(SAVE_DIR, filename)
+    
+    image.save(filepath)
+    return filepath
+
+@spaces.GPU(duration=60)
+def generate_image(
+    prompt: str,
+    seed: int,
+    randomize_seed: bool,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    num_inference_steps: int,
+    progress: gr.Progress = gr.Progress()
+):
+    try:
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+        generator = torch.Generator(device=device).manual_seed(seed)
+        
+        with torch.inference_mode():
+            image = pipe(
+                prompt=prompt,
+                width=width,
+                height=height,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            ).images[0]
+        
+        filepath = save_generated_image(image, prompt)
+        return image, seed
+        
+    except Exception as e:
+        raise gr.Error(f"Image generation failed: {str(e)}")
+    finally:
+        clear_memory()
+
+
+
+
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    position = gr.State(value="bottom-center")  # 여기로 이동
+    position = gr.State(value="bottom-center")
 
     gr.HTML("""
         <div class="main-title">
@@ -712,153 +770,173 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         </div>
     """)
 
-    with gr.Row(equal_height=True):
-        # 왼쪽 패널 (입력)
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="input-panel"):
-                input_image = gr.Image(
-                    type="pil",
-                    label="Upload Image",
-                    interactive=True,
-                    height=400
-                )
-                text_prompt = gr.Textbox(
-                    label="Object to Extract",
-                    placeholder="Enter what you want to extract...",
-                    interactive=True
-                )
-                with gr.Row():
-                    bg_prompt = gr.Textbox(
-                        label="Background Prompt (optional)",
-                        placeholder="Describe the background...",
-                        interactive=True,
-                        scale=3
-                    )
-                    aspect_ratio = gr.Dropdown(
-                        choices=["1:1", "16:9", "9:16", "4:3"],
-                        value="1:1",
-                        label="Aspect Ratio",
-                        interactive=True,
-                        visible=True,
-                        scale=1
-                    )
-
-                with gr.Group(elem_classes="controls-panel", visible=False) as object_controls:
-                    with gr.Column(scale=1):
-                        position = gr.State(value="bottom-center")  # 초기값 설정
-                        with gr.Row():
-                            btn_top_left = gr.Button("↖", elem_classes="position-btn")
-                            btn_top_center = gr.Button("↑", elem_classes="position-btn")
-                            btn_top_right = gr.Button("↗", elem_classes="position-btn")
-                        with gr.Row():
-                            btn_middle_left = gr.Button("←", elem_classes="position-btn")
-                            btn_middle_center = gr.Button("•", elem_classes="position-btn")
-                            btn_middle_right = gr.Button("→", elem_classes="position-btn")
-                        with gr.Row():
-                            btn_bottom_left = gr.Button("↙", elem_classes="position-btn")
-                            btn_bottom_center = gr.Button("↓", elem_classes="position-btn", value="selected")
-                            btn_bottom_right = gr.Button("↘", elem_classes="position-btn")
-                    with gr.Column(scale=1):
-                        scale_slider = gr.Slider(
-                            minimum=10,
-                            maximum=200,
-                            value=50,
-                            step=5,
-                            label="Object Size (%)"
+    with gr.Tabs():
+        with gr.Tab("Edit & Combine"):
+            with gr.Row(equal_height=True):
+                # 왼쪽 패널 (입력)
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="input-panel"):
+                        input_image = gr.Image(
+                            type="pil",
+                            label="Upload Image",
+                            interactive=True,
+                            height=400
                         )
-
-                process_btn = gr.Button(
-                    "Process",
-                    variant="primary",
-                    interactive=False,
-                    size="lg"
-                )
-
-        # 오른쪽 패널 (출력)
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="output-panel"):
-                with gr.Tab("Result"):
-                    combined_image = gr.Image(
-                        label="Combined Result",
-                        show_download_button=True,
-                        type="pil",
-                        height=400
-                    )
-                    
-                    # 텍스트 삽입 옵션을 Accordion으로 변경
-                    with gr.Accordion("Text Insertion Options", open=False):
-                        with gr.Group():
-                            with gr.Row():
-                                text_input = gr.Textbox(
-                                    label="Text Content",
-                                    placeholder="Enter text to add..."
-                                )
-                                text_position_type = gr.Radio(
-                                    choices=["Text Over Image", "Text Behind Image"],
-                                    value="Text Over Image",
-                                    label="Text Position"
+                        text_prompt = gr.Textbox(
+                            label="Object to Extract",
+                            placeholder="Enter what you want to extract...",
+                            interactive=True
+                        )
+                        with gr.Row():
+                            bg_prompt = gr.Textbox(
+                                label="Background Prompt (optional)",
+                                placeholder="Describe the background...",
+                                interactive=True,
+                                scale=3
+                            )
+                            aspect_ratio = gr.Dropdown(
+                                choices=["1:1", "16:9", "9:16", "4:3"],
+                                value="1:1",
+                                label="Aspect Ratio",
+                                interactive=True,
+                                visible=True,
+                                scale=1
+                            )
+
+                        with gr.Group(elem_classes="controls-panel", visible=False) as object_controls:
+                            with gr.Column(scale=1):
+                                position = gr.State(value="bottom-center")
+                                with gr.Row():
+                                    btn_top_left = gr.Button("↖", elem_classes="position-btn")
+                                    btn_top_center = gr.Button("↑", elem_classes="position-btn")
+                                    btn_top_right = gr.Button("↗", elem_classes="position-btn")
+                                with gr.Row():
+                                    btn_middle_left = gr.Button("←", elem_classes="position-btn")
+                                    btn_middle_center = gr.Button("•", elem_classes="position-btn")
+                                    btn_middle_right = gr.Button("→", elem_classes="position-btn")
+                                with gr.Row():
+                                    btn_bottom_left = gr.Button("↙", elem_classes="position-btn")
+                                    btn_bottom_center = gr.Button("↓", elem_classes="position-btn", value="selected")
+                                    btn_bottom_right = gr.Button("↘", elem_classes="position-btn")
+                            with gr.Column(scale=1):
+                                scale_slider = gr.Slider(
+                                    minimum=10,
+                                    maximum=200,
+                                    value=50,
+                                    step=5,
+                                    label="Object Size (%)"
                                 )
-                            
-                            with gr.Row():
-                                with gr.Column(scale=1):
-                                    font_choice = gr.Dropdown(
-                                        choices=["Default", "Korean Regular"],  # "Korean Son" 제거
-
-                                        value="Default",
-                                        label="Font Selection",
-                                        interactive=True
-                                    )
 
-                                    
-                                    font_size = gr.Slider(
-                                        minimum=10,
-                                        maximum=200,
-                                        value=40,
-                                        step=5,
-                                        label="Font Size"
-                                    )
-                                    color_dropdown = gr.Dropdown(
-                                        choices=["White", "Black", "Red", "Green", "Blue", "Yellow", "Purple"],
-                                        value="White",
-                                        label="Text Color"
-                                    )
-                                    thickness = gr.Slider(
-                                        minimum=0,
-                                        maximum=10,
-                                        value=1,
-                                        step=1,
-                                        label="Text Thickness"
-                                    )
-                                with gr.Column(scale=1):
-                                    opacity_slider = gr.Slider(
-                                        minimum=0,
-                                        maximum=255,
-                                        value=255,
-                                        step=1,
-                                        label="Opacity"
-                                    )
-                                    x_position = gr.Slider(
-                                        minimum=0,
-                                        maximum=100,
-                                        value=50,
-                                        step=1,
-                                        label="Left(0%)~Right(100%)"
+                        process_btn = gr.Button(
+                            "Process",
+                            variant="primary",
+                            interactive=False,
+                            size="lg"
+                        )
+
+                # 오른쪽 패널 (출력)
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="output-panel"):
+                        combined_image = gr.Image(
+                            label="Combined Result",
+                            show_download_button=True,
+                            type="pil",
+                            height=400
+                        )
+                        
+                        with gr.Accordion("Text Insertion Options", open=False):
+                            with gr.Group():
+                                with gr.Row():
+                                    text_input = gr.Textbox(
+                                        label="Text Content",
+                                        placeholder="Enter text to add..."
                                     )
-                                    y_position = gr.Slider(
-                                        minimum=0,
-                                        maximum=100,
-                                        value=50,
-                                        step=1,
-                                        label="High(0%)~Low(100%)"
+                                    text_position_type = gr.Radio(
+                                        choices=["Text Over Image", "Text Behind Image"],
+                                        value="Text Over Image",
+                                        label="Text Position"
                                     )
-                            add_text_btn = gr.Button("Apply Text", variant="primary")
+                                
+                                with gr.Row():
+                                    with gr.Column(scale=1):
+                                        font_choice = gr.Dropdown(
+                                            choices=["Default", "Korean Regular"],
+                                            value="Default",
+                                            label="Font Selection",
+                                            interactive=True
+                                        )
+                                        font_size = gr.Slider(
+                                            minimum=10,
+                                            maximum=200,
+                                            value=40,
+                                            step=5,
+                                            label="Font Size"
+                                        )
+                                        color_dropdown = gr.Dropdown(
+                                            choices=["White", "Black", "Red", "Green", "Blue", "Yellow", "Purple"],
+                                            value="White",
+                                            label="Text Color"
+                                        )
+                                        thickness = gr.Slider(
+                                            minimum=0,
+                                            maximum=10,
+                                            value=1,
+                                            step=1,
+                                            label="Text Thickness"
+                                        )
+                                    with gr.Column(scale=1):
+                                        opacity_slider = gr.Slider(
+                                            minimum=0,
+                                            maximum=255,
+                                            value=255,
+                                            step=1,
+                                            label="Opacity"
+                                        )
+                                        x_position = gr.Slider(
+                                            minimum=0,
+                                            maximum=100,
+                                            value=50,
+                                            step=1,
+                                            label="Left(0%)~Right(100%)"
+                                        )
+                                        y_position = gr.Slider(
+                                            minimum=0,
+                                            maximum=100,
+                                            value=50,
+                                            step=1,
+                                            label="High(0%)~Low(100%)"
+                                        )
+                                add_text_btn = gr.Button("Apply Text", variant="primary")
+
+                        extracted_image = gr.Image(
+                            label="Extracted Object",
+                            show_download_button=True,
+                            type="pil",
+                            height=200
+                        )
 
-                    extracted_image = gr.Image(
-                        label="Extracted Object",
-                        show_download_button=True,
-                        type="pil",
-                        height=200
-                    )
+        with gr.Tab("Generate Image"):
+            with gr.Column():
+                gen_prompt = gr.Textbox(
+                    label="Generation Prompt",
+                    placeholder="Enter your image generation prompt..."
+                )
+                with gr.Row():
+                    gen_width = gr.Slider(512, 1024, 768, step=64, label="Width")
+                    gen_height = gr.Slider(512, 1024, 768, step=64, label="Height")
+                
+                with gr.Row():
+                    guidance_scale = gr.Slider(1, 20, 7.5, step=0.5, label="Guidance Scale")
+                    num_steps = gr.Slider(1, 50, 30, step=1, label="Number of Steps")
+                
+                with gr.Row():
+                    seed = gr.Number(label="Seed", value=-1)
+                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                
+                generate_btn = gr.Button("Generate Image", variant="primary")
+                
+                output_image = gr.Image(label="Generated Image", type="pil")
+                output_seed = gr.Number(label="Used Seed", interactive=False)
 
     # CSS 클래스를 위한 스타일 추가
     gr.HTML("""
@@ -870,7 +948,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         </style>
     """)
 
-# 버튼 클릭 이벤트 바인딩
+    # 버튼 클릭 이벤트 바인딩
     position_mapping = {
         btn_top_left: "top-left",
         btn_top_center: "top-center",
@@ -885,10 +963,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
 
     for btn, pos in position_mapping.items():
         btn.click(
-            fn=lambda pos=pos: update_position(pos),  # 클로저 문제 해결을 위해 수정
+            fn=lambda pos=pos: update_position(pos),
             outputs=position
         )
- 
 
     # 이벤트 바인딩
     bg_prompt.change(
@@ -926,12 +1003,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         queue=True
     )
 
-# 이벤트 바인딩 부분에서
     add_text_btn.click(
         fn=add_text_to_image,
         inputs=[
-            combined_image,  # 첫 번째 인자로 이미지
-            text_input,      # 두 번째 인자로 텍스트
+            combined_image,
+            text_input,
             font_size,
             color_dropdown,
             opacity_slider,
@@ -942,10 +1018,22 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             font_choice
         ],
         outputs=combined_image,
-        api_name="add_text"  # API 이름 추가
+        api_name="add_text"
     )
-    
 
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[
+            gen_prompt,
+            seed,
+            randomize_seed,
+            gen_width,
+            gen_height,
+            guidance_scale,
+            num_steps,
+        ],
+        outputs=[output_image, output_seed]
+    )
 
 demo.queue(max_size=5)
 demo.launch(