Add CPU, GPU warm up process when app load first time

app.py

@@ -189,6 +189,114 @@ print("\n" + "=" * 60)
 print("All models loaded successfully!")
 print("=" * 60 + "\n")
 
+# Warm-up: 첫 번째 추론 지연 감소를 위한 모델 초기화
+# JIT 컴파일, CUDA 커널 로딩 등을 미리 수행
+print("=" * 60)
+print("Warming up models (CPU)...")
+print("=" * 60)
+
+def warmup_models_cpu():
+    """앱 시작 시 CPU 모델 초기화를 위한 Warm-up 함수"""
+    try:
+        # CPU에서 텍스트 임베딩 Warm-up (Tokenizer + Text Encoder 초기화)
+        print("[CPU Warm-up 1/2] Text Encoder warm-up...")
+        with torch.no_grad():
+            dummy_prompt = "a photo of clothing"
+            dummy_tokens = tokenizer_one(
+                dummy_prompt,
+                padding="max_length",
+                max_length=tokenizer_one.model_max_length,
+                truncation=True,
+                return_tensors="pt"
+            )
+            # CPU에서 실행 가능한 초기화
+            _ = text_encoder_one(dummy_tokens.input_ids, output_hidden_states=True)
+        print("✓ Text Encoder warmed up")
+        
+        # Tensor 변환 Warm-up
+        print("[CPU Warm-up 2/2] Tensor transform warm-up...")
+        dummy_img = Image.new('RGB', (768, 1024), color='white')
+        _ = tensor_transfrom(dummy_img)
+        print("✓ Tensor transform warmed up")
+        
+        return True
+    except Exception as e:
+        print(f"⚠ CPU Warm-up partially completed: {e}")
+        return False
+
+# CPU Warm-up 실행
+warmup_success = warmup_models_cpu()
+if warmup_success:
+    print("\n✓ CPU warm-up completed successfully")
+else:
+    print("\n⚠ CPU warm-up completed with warnings")
+print("=" * 60 + "\n")
+
+# GPU Warm-up 함수 (앱 로드 시 자동 실행)
+# torch.compile() 첫 번째 컴파일을 미리 수행
+@spaces.GPU
+def warmup_gpu():
+    """앱 로드 시 GPU 모델 초기화를 위한 Warm-up 함수 (torch.compile 첫 호출)"""
+    try:
+        device = "cuda"
+        print("=" * 60)
+        print("GPU Warm-up: Triggering torch.compile() first compilation...")
+        print("=" * 60)
+        
+        # 모델을 GPU로 이동
+        pipe.to(device)
+        pipe.unet_encoder.to(device)
+        
+        # 더미 텐서 생성
+        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...")
+                dummy_prompt = "a photo of white t-shirt"
+                (
+                    prompt_embeds,
+                    negative_prompt_embeds,
+                    pooled_prompt_embeds,
+                    negative_pooled_prompt_embeds,
+                ) = pipe.encode_prompt(
+                    dummy_prompt,
+                    num_images_per_prompt=1,
+                    do_classifier_free_guidance=True,
+                    negative_prompt="low quality",
+                )
+                print("✓ Text Encoder GPU warmed up")
+                
+                # 2. 더미 이미지로 VAE 인코딩 (VAE GPU warm-up)
+                print("[GPU Warm-up 2/3] 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")
+                
+                # 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)")
+        
+        # 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("=" * 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.")
+        return f"GPU Warm-up skipped: {e}"
+
 
 # 이미지 전처리 함수
 def preprocess_image(image):
@@ -628,9 +736,19 @@ with image_blocks as demo:
         api_name='tryon'
     )
     
+    # GPU Warm-up 상태 표시용 (숨김)
+    warmup_status = gr.Textbox(visible=False)
+    
+    # 앱 로드 시 GPU Warm-up 자동 실행 (torch.compile 첫 컴파일)
+    demo.load(
+        fn=warmup_gpu,
+        inputs=None,
+        outputs=warmup_status,
+    )
     
     print("✓ Gradio interface components created")
     print("✓ Event handlers configured")
+    print("✓ GPU warm-up scheduled on app load")
 
 print("\n" + "=" * 60)
 print("Gradio Application Interface Created Successfully!")