Warm up : model move to GPU when inference at first time

app.py

@@ -232,33 +232,39 @@ else:
     print("\n⚠ CPU warm-up completed with warnings")
 print("=" * 60 + "\n")
 
+# torch.compile 오류 시 eager 모드로 폴백 설정
+# 커스텀 UNet forward 메서드 호환성 문제 대응
+try:
+    import torch._dynamo
+    torch._dynamo.config.suppress_errors = True
+    print("✓ torch._dynamo.config.suppress_errors enabled (fallback to eager mode on error)")
+except Exception as e:
+    print(f"⚠ torch._dynamo config not available: {e}")
+
 # GPU Warm-up 함수 (앱 로드 시 자동 실행)
-# torch.compile() 첫 번째 컴파일을 미리 수행
+# Text Encoder, VAE GPU 로딩 및 CUDA 커널 초기화
 @spaces.GPU
 def warmup_gpu():
-    """앱 로드 시 GPU 모델 초기화를 위한 Warm-up 함수 (torch.compile 첫 호출)"""
+    """앱 로드 시 GPU 모델 초기화를 위한 Warm-up 함수"""
     try:
         device = "cuda"
         print("=" * 60)
-        print("GPU Warm-up: Triggering torch.compile() first compilation...")
+        print("GPU Warm-up: Loading models to GPU and initializing CUDA kernels...")
         print("=" * 60)
         
         # 모델을 GPU로 이동
+        print("[GPU Warm-up 1/4] Moving models to GPU...")
         pipe.to(device)
         pipe.unet_encoder.to(device)
+        print("✓ Models moved to GPU")
         
         # 더미 텐서 생성
         with torch.no_grad():
             with torch.cuda.amp.autocast():
                 # 1. 더미 프롬프트 임베딩 생성 (Text Encoder GPU warm-up)
-                print("[GPU Warm-up 1/3] Text Encoder GPU warm-up...")
+                print("[GPU Warm-up 2/4] Text Encoder GPU warm-up...")
                 dummy_prompt = "a photo of white t-shirt"
-                (
-                    prompt_embeds,
-                    negative_prompt_embeds,
-                    pooled_prompt_embeds,
-                    negative_pooled_prompt_embeds,
-                ) = pipe.encode_prompt(
+                _ = pipe.encode_prompt(
                     dummy_prompt,
                     num_images_per_prompt=1,
                     do_classifier_free_guidance=True,
@@ -266,35 +272,31 @@ def warmup_gpu():
                 )
                 print("✓ Text Encoder GPU warmed up")
                 
-                # 2. 더미 이미지로 VAE 인코딩 (VAE GPU warm-up)
-                print("[GPU Warm-up 2/3] VAE GPU warm-up...")
+                # 2. 더미 이미지로 VAE 인코딩/디코딩 (VAE GPU warm-up)
+                print("[GPU Warm-up 3/4] VAE GPU warm-up...")
                 dummy_img = torch.randn(1, 3, 1024, 768).to(device, torch.float16)
-                _ = pipe.vae.encode(dummy_img)
-                print("✓ VAE GPU warmed up")
+                latents = pipe.vae.encode(dummy_img).latent_dist.sample()
+                _ = pipe.vae.decode(latents)
+                print("✓ VAE GPU warmed up (encode + decode)")
                 
-                # 3. UNet 간단한 forward pass (UNet + torch.compile warm-up)
-                print("[GPU Warm-up 3/3] UNet GPU warm-up (torch.compile trigger)...")
-                dummy_latent = torch.randn(1, 4, 128, 96).to(device, torch.float16)
-                dummy_timestep = torch.tensor([999]).to(device)
-                _ = pipe.unet(
-                    dummy_latent,
-                    dummy_timestep,
-                    encoder_hidden_states=prompt_embeds.to(device, torch.float16),
-                )
-                print("✓ UNet GPU warmed up (torch.compile triggered)")
+                # 3. CUDA 동기화 (커널 로딩 완료 대기)
+                print("[GPU Warm-up 4/4] CUDA synchronization...")
+                torch.cuda.synchronize()
+                print("✓ CUDA kernels initialized")
         
         # GPU 메모리 정리
         torch.cuda.empty_cache()
         
         print("\n" + "=" * 60)
-        print("✓ GPU Warm-up completed! torch.compile() compilation done.")
-        print("  All subsequent requests will be faster.")
+        print("✓ GPU Warm-up completed!")
+        print("  Text Encoder, VAE ready. UNet will compile on first request.")
+        print("  (torch.compile errors will fallback to eager mode)")
         print("=" * 60 + "\n")
         
         return "GPU Warm-up completed successfully!"
     except Exception as e:
         print(f"\n⚠ GPU Warm-up failed: {e}")
-        print("  First user request will trigger compilation instead.")
+        print("  Models will be loaded on first user request.")
         return f"GPU Warm-up skipped: {e}"
 
 
@@ -662,8 +664,9 @@ with image_blocks as demo:
     print("✓ Gradio Blocks created")
     
     gr.Markdown("## DXCO : GENAI-VTON")
-    gr.Markdown("임성남, 윤지영, 조민주 based on IDM-VTON")
-    gr.Markdown("이미지는 3:4비율(384x512 또는 768x1024)로 올려주세요")
+    gr.Markdown("임성남, 윤지영, 조민주 based on IDM-VTON")    
+    gr.Markdown("* 맨 처음 추론 시 [5분] 걸림 - compile과 GPU warm-up *")
+    gr.Markdown("권장 이미지 사이즈 - 3:4비율(384x512,768x1024)")
     
     with gr.Row():
         with gr.Column():