Update app.py

app.py

@@ -17,7 +17,7 @@ initialization_message = "모델 로딩 중... 잠시만 기다려주세요."
 # 간단한 인용 정보 추가
 _CITE_ = """PuLID: Person-under-Language Image Diffusion Model"""
 
-# GPU 사용 가능 여부 확인 및 장치 설정
+# GPU 사용 가능 여부 확인 및 장치 설정 - 메인 프로세스에서는 호출하지 않음
 def get_device():
     if torch.cuda.is_available():
         return torch.device('cuda')
@@ -44,8 +44,8 @@ def get_models(name: str, device, offload: bool):
 
 class FluxGenerator:
     def __init__(self):
-        # GPU 사용 가능 여부에 따라 장치 설정
-        self.device = get_device()
+        # GPU 초기화는 Spaces GPU 데코레이터 안에서만 수행
+        self.device = None  # 초기화 시점에는 device를 할당하지 않음
         self.offload = False
         self.model_name = 'flux-dev'
         self.initialized = False
@@ -63,6 +63,9 @@ class FluxGenerator:
             from pulid.pipeline_flux import PuLIDPipeline
             from flux.sampling import prepare
             
+            # 이 시점에서 장치 설정 (GPU 데코레이터 내에서만 호출됨)
+            self.device = get_device()
+            
             print("모델 초기화 시작...")
             self.model, self.ae, self.t5, self.clip_model = get_models(
                 self.model_name,
@@ -98,11 +101,12 @@ class FluxGenerator:
             initialization_message = f"모델 로딩 실패: {str(e)}"
 
 
-# 지연 로딩을 위한 백그라운드 초기화 함수
-def initialize_models_in_background():
+# 지연 로딩을 위한 백그라운드 초기화 함수 - GPU 데코레이터로 변경
+@spaces.GPU(duration=60)
+def initialize_models():
     global flux_generator, model_initialized, initialization_message
     
-    print("백그라운드에서 모델 초기화 시작...")
+    print("GPU 데코레이터 내에서 모델 초기화 시작...")
     
     try:
         # 지연 임포트
@@ -112,19 +116,18 @@ def initialize_models_in_background():
         
         # 모델 초기화
         flux_generator = FluxGenerator()
-        
-        # 30초 후에 초기화 시작 (UI가 먼저 로드되도록)
-        time.sleep(30)
         flux_generator.initialize()
         
         model_initialized = flux_generator.initialized
         
     except Exception as e:
         import traceback
-        error_msg = f"백그라운드 초기화 중 오류 발생: {str(e)}\n{traceback.format_exc()}"
+        error_msg = f"초기화 중 오류 발생: {str(e)}\n{traceback.format_exc()}"
         print(error_msg)
         model_initialized = False
         initialization_message = f"모델 초기화 오류: {str(e)}"
+    
+    return initialization_message
 
 
 # 모델 상태 확인 함수
@@ -151,7 +154,7 @@ def generate_image(
     
     # 모델이 초기화되지 않았으면 오류 메시지 반환
     if not model_initialized:
-        return None, "모델 초기화가 완료되지 않았습니다. 잠시 후 다시 시도해주세요."
+        return None, "모델 초기화가 완료되지 않았습니다. 모델 초기화 버튼을 눌러주세요."
     
     # ID 이미지가 없으면 실행 불가
     if id_image is None:
@@ -322,6 +325,11 @@ def create_demo():
         
         # 모델 상태 표시
         status_box = gr.Textbox(label="모델 상태", value=initialization_message)
+        
+        # 초기화 버튼 추가 (백그라운드 초기화 대신 명시적 초기화 버튼 사용)
+        init_btn = gr.Button("모델 초기화")
+        init_btn.click(fn=initialize_models, inputs=[], outputs=[status_box])
+        
         refresh_btn = gr.Button("상태 새로고침")
         refresh_btn.click(fn=check_model_status, inputs=[], outputs=[status_box])
             
@@ -369,9 +377,6 @@ def create_demo():
                        id_weight, neg_prompt, true_cfg, gamma, eta]
             )
 
-        # 주기적 상태 업데이트 설정
-        demo.load(fn=check_model_status, inputs=[], outputs=[status_box], every=5)  # 5초마다 업데이트
-
         # Gradio 이벤트 연결
         generate_btn.click(
             fn=generate_image,
@@ -396,18 +401,9 @@ if __name__ == "__main__":
 
     print("Hugging Face Spaces 환경에서 실행 중입니다. GPU 할당을 요청합니다.")
     
-    # UI가 먼저 로드되도록 백그라운드에서 모델 초기화 시작
-    threading.Thread(target=initialize_models_in_background, daemon=True).start()
-    
-    # 메모리 사용량 제한 설정 (환경에 따라 조정 필요)
-    try:
-        import torch.cuda
-        if torch.cuda.is_available():
-            # GPU 메모리 사용량 제한 (필요시 조정)
-            torch.cuda.set_per_process_memory_fraction(0.8)  # 최대 80% GPU 메모리 사용
-    except Exception as e:
-        print(f"메모리 제한 설정 중 오류: {e}")
+    # 메인 프로세스에서는 CUDA 초기화하지 않음
+    # 백그라운드 스레드 대신 명시적 버튼으로 초기화
     
     demo = create_demo()
     # 디버그 모드 활성화
-    demo.queue().launch(server_name="0.0.0.0", server_port=args.port, debug=True)
+    demo.queue().launch(server_name="0.0.0.0", server_port=args.port)