Update app.py

app.py

@@ -2,9 +2,11 @@ import spaces
 import argparse
 import os
 import time
+import gc
 from os import path
 import shutil
 from datetime import datetime
+import traceback
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 import gradio as gr
@@ -20,7 +22,9 @@ os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
+# GPU 메모리 설정 최적화
 torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.benchmark = True  # 반복적인 크기의 입력에 대해 성능 향상
 
 def filter_prompt(prompt):
     # 부적절한 키워드 목록
@@ -53,17 +57,41 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
-# Model initialization
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=True)
+# 글로벌 변수로 파이프라인 선언
+pipe = None
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
-pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device="cuda", dtype=torch.bfloat16)
-
-# Add safety checker
-pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+# 모델 초기화 함수 (지연 로딩)
+def initialize_model():
+    global pipe
+    
+    # 이미 로드된 경우 다시 로드하지 않음
+    if pipe is not None:
+        return
+    
+    try:
+        if not path.exists(cache_path):
+            os.makedirs(cache_path, exist_ok=True)
+        
+        # 메모리 확보를 위한 가비지 컬렉션 실행
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        with timer("모델 로딩"):
+            pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+            lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
+            pipe.load_lora_weights(lora_path)
+            pipe.fuse_lora(lora_scale=0.125)
+            pipe.to(device="cuda", dtype=torch.bfloat16)
+            
+            # 안전 검사기 추가
+            pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+            
+        print("모델 로딩 완료")
+        return True
+    except Exception as e:
+        print(f"모델 로딩 중 오류 발생: {str(e)}")
+        traceback.print_exc()
+        return False
 
 css = """
 footer {display: none !important}
@@ -106,6 +134,20 @@ footer {display: none !important}
     width: 100% !important;
     max-width: 100% !important;
 }
+.loading-indicator {
+    text-align: center;
+    padding: 20px;
+    font-weight: bold;
+    color: #4B79A1;
+}
+.error-message {
+    background-color: rgba(255, 0, 0, 0.1);
+    color: red;
+    padding: 10px;
+    border-radius: 8px;
+    margin-top: 10px;
+    text-align: center;
+}
 """
 
 # Create Gradio interface
@@ -119,6 +161,10 @@ 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"])
+
     with gr.Row():
         with gr.Column(scale=3):
             prompt = gr.Textbox(
@@ -161,7 +207,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     )
                 
                 def get_random_seed():
-                    return torch.randint(0, 1000000, (1,)).item()
+                    return int(torch.randint(0, 1000000, (1,)).item())
                 
                 seed = gr.Number(
                     label="Seed (random by default, set for reproducibility)",
@@ -211,34 +257,92 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     
     @spaces.GPU
     def process_image(height, width, steps, scales, prompt, seed):
+        # 모델 초기화 상태 확인
+        if pipe is None:
+            loading_status.update("모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다.", visible=True)
+            
+            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)
+        
+        # 입력값 검증
+        if not prompt or prompt.strip() == "":
+            error_message.update("이미지 설명을 입력해주세요.", visible=True)
+            return None
+            
         # 프롬프트 필터링
         is_safe, filtered_prompt = filter_prompt(prompt)
         if not is_safe:
-            gr.Warning("부적절한 내용이 ��함된 프롬프트입니다.")
+            error_message.update("부적절한 내용이 포함된 프롬프트입니다.", visible=True)
             return None
             
-        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-            try:
+        # 에러 메시지 초기화
+        error_message.update(visible=False)
+        loading_status.update("이미지를 생성 중입니다...", visible=True)
+            
+        try:
+            # 메모리 확보를 위한 가비지 콜렉션
+            gc.collect()
+            torch.cuda.empty_cache()
+            
+            # 시드 값 확인 및 보정
+            if seed is None or not isinstance(seed, (int, float)):
+                seed = get_random_seed()
+            else:
+                seed = int(seed)  # 타입 변환 안전하게 처리
+                
+            # 이미지 생성
+            with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+                generator = torch.Generator(device="cuda").manual_seed(seed)
+                
+                # 높이와 너비를 64의 배수로 조정 (FLUX 모델 요구사항)
+                height = (int(height) // 64) * 64
+                width = (int(width) // 64) * 64
+                
+                # 안전장치 - 최대값 제한
+                steps = min(int(steps), 25)
+                scales = max(min(float(scales), 5.0), 0.0)
+                
                 generated_image = pipe(
                     prompt=[filtered_prompt],
-                    generator=torch.Generator().manual_seed(int(seed)),
-                    num_inference_steps=int(steps),
-                    guidance_scale=float(scales),
-                    height=int(height),
-                    width=int(width),
+                    generator=generator,
+                    num_inference_steps=steps,
+                    guidance_scale=scales,
+                    height=height,
+                    width=width,
                     max_sequence_length=256
                 ).images[0]
                 
+                loading_status.update(visible=False)
                 return generated_image
-            except Exception as e:
-                print(f"Error in image generation: {str(e)}")
-                return None
+                
+        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
     
     def update_seed():
         return get_random_seed()
+        
+    # 버튼 클릭 이벤트 - 모든 UI 요소 초기화 추가
+    def on_generate_click(height, width, steps, scales, prompt, seed):
+        error_message.update(visible=False)
+        return process_image(height, width, steps, scales, prompt, seed)
 
     generate_btn.click(
-        process_image,
+        on_generate_click,
         inputs=[height, width, steps, scales, prompt, seed],
         outputs=[output]
     )
@@ -247,11 +351,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         update_seed,
         outputs=[seed]
     )
-    
-    generate_btn.click(
-        update_seed,
-        outputs=[seed]
-    )
 
 if __name__ == "__main__":
-    demo.launch()
+    # 앱 시작 시 모델 미리 로드하지 않음 (첫 요청 시 지연 로딩)
+    demo.queue(concurrency_count=1, max_size=10).launch()