Update app.py

app.py

@@ -66,7 +66,7 @@ def initialize_model():
     
     # 이미 로드된 경우 다시 로드하지 않음
     if pipe is not None:
-        return
+        return True
     
     try:
         if not path.exists(cache_path):
@@ -161,9 +161,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         </div>
     """)
 
-    # 상태 표시 변수
-    error_message = gr.HTML(visible=False, elem_classes=["error-message"])
-    loading_status = gr.HTML(visible=False, elem_classes=["loading-indicator"])
+    # 상태 표시 변수 (HTML 대신 Textbox 사용)
+    error_message = gr.Textbox(
+        value="", 
+        label="Error", 
+        visible=False, 
+        elem_classes=["error-message"]
+    )
+    
+    loading_status = gr.Textbox(
+        value="", 
+        label="Status", 
+        visible=False, 
+        elem_classes=["loading-indicator"]
+    )
 
     with gr.Row():
         with gr.Column(scale=3):
@@ -257,32 +268,24 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     
     @spaces.GPU
     def process_image(height, width, steps, scales, prompt, seed):
+        global pipe
+        
         # 모델 초기화 상태 확인
         if pipe is None:
-            loading_status.update("모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다.", visible=True)
+            return None, "모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다.", True, "", False
             
             model_loaded = initialize_model()
             if not model_loaded:
-                error_message.update("모델 로딩 중 오류가 발생했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", visible=True)
-                loading_status.update(visible=False)
-                return None
-                
-            loading_status.update(visible=False)
+                return None, "", False, "모델 로딩 중 오류가 발생했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True
         
         # 입력값 검증
         if not prompt or prompt.strip() == "":
-            error_message.update("이미지 설명을 입력해주세요.", visible=True)
-            return None
+            return None, "", False, "이미지 설명을 입력해주세요.", True
             
         # 프롬프트 필터링
         is_safe, filtered_prompt = filter_prompt(prompt)
         if not is_safe:
-            error_message.update("부적절한 내용이 포함된 프롬프트입니다.", visible=True)
-            return None
-            
-        # 에러 메시지 초기화
-        error_message.update(visible=False)
-        loading_status.update("이미지를 생성 중입니다...", visible=True)
+            return None, "", False, "부적절한 내용이 포함된 프롬프트입니다.", True
             
         try:
             # 메모리 확보를 위한 가비지 콜렉션
@@ -295,6 +298,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             else:
                 seed = int(seed)  # 타입 변환 안전하게 처리
                 
+            # 이미지 생성 상태 메시지
+            loading_message = "이미지를 생성 중입니다..."
+                
             # 이미지 생성
             with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
                 generator = torch.Generator(device="cuda").manual_seed(seed)
@@ -317,42 +323,46 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     max_sequence_length=256
                 ).images[0]
                 
-                loading_status.update(visible=False)
-                return generated_image
+                # 성공 시 이미지 반환, 상태 메시지 숨김
+                return generated_image, "", False, "", False
                 
         except Exception as e:
             error_msg = f"이미지 생성 중 오류가 발생했습니다: {str(e)}"
             print(error_msg)
             traceback.print_exc()
-            error_message.update(error_msg, visible=True)
-            loading_status.update(visible=False)
             
             # 오류 후 메모리 정리
             gc.collect()
             torch.cuda.empty_cache()
             
-            return None
+            return None, "", False, error_msg, True
     
     def update_seed():
         return get_random_seed()
-        
-    # 버튼 클릭 이벤트 - 모든 UI 요소 초기화 추가
-    def on_generate_click(height, width, steps, scales, prompt, seed):
-        error_message.update(visible=False)
+    
+    # 이미지 생성 준비 함수
+    def prepare_generation(height, width, steps, scales, prompt, seed):
+        # 모델이 아직 로드되지 않았다면 로드
+        if pipe is None:
+            is_loaded = initialize_model()
+            if not is_loaded:
+                return None, "모델 로딩에 실패했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True, "", False
+            
+        # 생성 프로세스 시작
         return process_image(height, width, steps, scales, prompt, seed)
 
+    # 버튼 클릭 이벤트 연결
     generate_btn.click(
-        on_generate_click,
+        fn=prepare_generation,
         inputs=[height, width, steps, scales, prompt, seed],
-        outputs=[output]
+        outputs=[output, loading_status, loading_status, error_message, error_message]
     )
     
     randomize_seed.click(
-        update_seed,
+        fn=update_seed,
         outputs=[seed]
     )
 
 if __name__ == "__main__":
     # 앱 시작 시 모델 미리 로드하지 않음 (첫 요청 시 지연 로딩)
-    demo.queue(max_size=10).launch()
-    
+    demo.queue(max_size=10).launch()