Update app.py

app.py

@@ -304,11 +304,15 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
             new_size = (int(img.width * ratio), int(img.height * ratio))
             img = img.resize(new_size, Image.LANCZOS)
 
-        # 메모리 사용량 모니터링
-        if torch.cuda.is_available():
-            torch.cuda.reset_peak_memory_stats()
+        # CUDA 메모리 관리 수정
+        try:
+            if torch.cuda.is_available():
+                current_device = torch.cuda.current_device()
+                torch.cuda.reset_peak_memory_stats(current_device)
+        except Exception as e:
+            print(f"CUDA memory stats reset failed: {e}")
 
-        with torch.cuda.amp.autocast():
+        with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
             mask, bbox, time_log = _gpu_process(img, prompt)
             masked_alpha = apply_mask(img, mask, defringe=True)
 
@@ -318,11 +322,12 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
         else:
             combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
 
-        clear_memory()  # 중간 메모리 정리
+        clear_memory()
 
         return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)
     except Exception as e:
         clear_memory()
+        print(f"Processing error: {str(e)}")
         raise gr.Error(f"Processing failed: {str(e)}")
 
 def on_change_bbox(prompts: dict[str, Any] | None):
@@ -621,69 +626,66 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
 
 # UI 구성 부분의 마지막에서 예시 테이블 부분을 수정
 
-    # 예시 테이블 추가
-    examples_path = "examples"  # 예시 이미지가 있는 디렉토리
-    gr.HTML(f"""
-        <div style="margin-top: 2em; padding: 1em;">
-            <h2 style="text-align: center; color: #2196F3; margin-bottom: 1em;">Examples</h2>
-            <div class="examples-container" style="overflow-x: auto;">
-                <table style="width: 100%; border-collapse: collapse; margin: auto; background: white; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
-                    <thead>
-                        <tr style="background: #f7f7f7;">
-                            <th style="padding: 1em; border: 1px solid #ddd; min-width: 120px;">Step</th>
-                            <th style="padding: 1em; border: 1px solid #ddd; min-width: 200px;">Example 1</th>
-                            <th style="padding: 1em; border: 1px solid #ddd; min-width: 200px;">Example 2</th>
-                        </tr>
-                    </thead>
-                    <tbody>
-                        <tr>
-                            <td style="padding: 1em; border: 1px solid #ddd;">Upload Image</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/aa1.png" style="max-width: 150px; height: auto;" alt="Example 1 Input">
-                            </td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/ab1.png" style="max-width: 150px; height: auto;" alt="Example 2 Input">
-                            </td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 1em; border: 1px solid #ddd;">Object Extract</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">text</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">sunglasses</td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 1em; border: 1px solid #ddd;">Background Prompt</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">The Alps in Spring with Blossoming Flowers</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">하와이 beach</td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 1em; border: 1px solid #ddd;">Object Cutoff Image</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/aa2.png" style="max-width: 150px; height: auto;" alt="Example 1 Cutoff">
-                            </td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/ab2.png" style="max-width: 150px; height: auto;" alt="Example 2 Cutoff">
-                            </td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 1em; border: 1px solid #ddd;">Result Image</td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/aa3.png" style="max-width: 150px; height: auto;" alt="Example 1 Result">
-                            </td>
-                            <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
-                                <img src="examples/ab3.png" style="max-width: 150px; height: auto;" alt="Example 2 Result">
-                            </td>
-                        </tr>
-                    </tbody>
-                </table>
-            </div>
+# 예시 테이블 HTML 수정
+gr.HTML("""
+    <div style="margin-top: 2em; padding: 1em;">
+        <h2 style="text-align: center; color: #2196F3; margin-bottom: 1em;">Examples</h2>
+        <div class="examples-container" style="overflow-x: auto;">
+            <table style="width: 100%; border-collapse: collapse; margin: auto; background: white; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
+                <thead>
+                    <tr style="background: #f7f7f7;">
+                        <th style="padding: 1em; border: 1px solid #ddd; min-width: 120px;">Step</th>
+                        <th style="padding: 1em; border: 1px solid #ddd; min-width: 200px;">Example 1</th>
+                        <th style="padding: 1em; border: 1px solid #ddd; min-width: 200px;">Example 2</th>
+                    </tr>
+                </thead>
+                <tbody>
+                    <tr>
+                        <td style="padding: 1em; border: 1px solid #ddd;">Upload Image</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example1_input.jpg" style="max-width: 150px; height: auto;" alt="Example 1 Input">
+                        </td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example2_input.jpg" style="max-width: 150px; height: auto;" alt="Example 2 Input">
+                        </td>
+                    </tr>
+                    <tr>
+                        <td style="padding: 1em; border: 1px solid #ddd;">Object Extract</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">text</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">sunglasses</td>
+                    </tr>
+                    <tr>
+                        <td style="padding: 1em; border: 1px solid #ddd;">Background Prompt</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">The Alps in Spring with Blossoming Flowers</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">하와이 beach</td>
+                    </tr>
+                    <tr>
+                        <td style="padding: 1em; border: 1px solid #ddd;">Object Cutoff Image</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example1_cutoff.jpg" style="max-width: 150px; height: auto;" alt="Example 1 Cutoff">
+                        </td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example2_cutoff.jpg" style="max-width: 150px; height: auto;" alt="Example 2 Cutoff">
+                        </td>
+                    </tr>
+                    <tr>
+                        <td style="padding: 1em; border: 1px solid #ddd;">Result Image</td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example1_result.jpg" style="max-width: 150px; height: auto;" alt="Example 1 Result">
+                        </td>
+                        <td style="padding: 1em; border: 1px solid #ddd; text-align: center;">
+                            <img src="https://raw.githubusercontent.com/your-username/your-repo/main/examples/example2_result.jpg" style="max-width: 150px; height: auto;" alt="Example 2 Result">
+                        </td>
+                    </tr>
+                </tbody>
+            </table>
         </div>
-    """)
-    
-demo.queue(max_size=10)  # 큐 크기 제한
+    </div>
+""")
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
-    share=False,
-    max_threads=2,  # 스레드 수 제한
-    allowed_paths=["examples"]
+    share=True,  # 공개 링크 활성화
+    max_threads=2,
+    allowed_paths=[examples_path]  # 수정된 예제 경로 사용
 )