Support zeroGPU env

app.py

@@ -30,6 +30,10 @@ from torchvision.transforms.functional import to_pil_image
 # import pillow_heif  # HEIC 이미지 처리용 (아이폰 촬영 사진 포맷)
 from urllib.parse import urlparse
 
+
+# zeroGPU 환경에서 compile 사용 여부
+is_compile_for_zeroGPU = True # True: compile 사용, False: compile 사용 안 함
+
 # SSL 경고 억제
 warnings.filterwarnings("ignore", message=".*OpenSSL.*")
 warnings.filterwarnings("ignore", category=UserWarning, module="urllib3")
@@ -68,14 +72,17 @@ unet = UNet2DConditionModel.from_pretrained(
 unet.requires_grad_(False)
 # torch.compile() 적용 - 추론 속도 20-40% 향상 (PyTorch 2.0+)
 # 주의: 첫 번째 추론은 컴파일로 인해 느릴 수 있음
-if hasattr(torch, 'compile'):
-    try:
-        unet = torch.compile(unet, mode="reduce-overhead")
-        print("✓ UNet model loaded and compiled successfully")
-    except Exception as e:
-        print(f"✓ UNet model loaded (compile skipped: {e})")
-else:
+if is_compile_for_zeroGPU == True:
     print("✓ UNet model loaded successfully")
+else:
+    if hasattr(torch, 'compile'):
+        try:
+            unet = torch.compile(unet, mode="reduce-overhead")
+            print("✓ UNet model loaded and compiled successfully")
+        except Exception as e:
+            print(f"✓ UNet model loaded (compile skipped: {e})")
+    else:
+        print("✓ UNet model loaded successfully")
 
 print("\n[2/10] Loading tokenizers...")
 tokenizer_one = AutoTokenizer.from_pretrained(
@@ -123,14 +130,17 @@ vae = AutoencoderKL.from_pretrained(base_path,
                                     torch_dtype=torch.float16,
 )
 # torch.compile() 적용 - VAE 인코딩/디코딩 속도 향상
-if hasattr(torch, 'compile'):
-    try:
-        vae = torch.compile(vae, mode="reduce-overhead")
-        print("✓ VAE loaded and compiled successfully")
-    except Exception as e:
-        print(f"✓ VAE loaded (compile skipped: {e})")
-else:
+if is_compile_for_zeroGPU == True:
     print("✓ VAE loaded successfully")
+else:
+    if hasattr(torch, 'compile'):
+        try:
+            vae = torch.compile(vae, mode="reduce-overhead")
+            print("✓ VAE loaded and compiled successfully")
+        except Exception as e:
+            print(f"✓ VAE loaded (compile skipped: {e})")
+    else:
+        print("✓ VAE loaded successfully")
 
 print("\n[7/10] Loading UNet Encoder...")
 UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
@@ -139,14 +149,17 @@ UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
     torch_dtype=torch.float16,
 )
 # torch.compile() 적용 - UNet Encoder 속도 향상
-if hasattr(torch, 'compile'):
-    try:
-        UNet_Encoder = torch.compile(UNet_Encoder, mode="reduce-overhead")
-        print("✓ UNet Encoder loaded and compiled successfully")
-    except Exception as e:
-        print(f"✓ UNet Encoder loaded (compile skipped: {e})")
-else:
+if is_compile_for_zeroGPU == True:
     print("✓ UNet Encoder loaded successfully")
+else:
+    if hasattr(torch, 'compile'):
+        try:
+            UNet_Encoder = torch.compile(UNet_Encoder, mode="reduce-overhead")
+            print("✓ UNet Encoder loaded and compiled successfully")
+        except Exception as e:
+            print(f"✓ UNet Encoder loaded (compile skipped: {e})")
+    else:
+        print("✓ UNet Encoder loaded successfully")
 
 print("\n[8/10] Initializing parsing and openpose models...")
 parsing_model = Parsing(0)
@@ -243,12 +256,15 @@ 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}")
+if is_compile_for_zeroGPU == True:
+    print("✓ torch.compile is disabled for ZeroGPU")
+else:
+    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 함수 (앱 로드 시 자동 실행)
 # Text Encoder, VAE GPU 로딩 및 CUDA 커널 초기화
@@ -752,11 +768,14 @@ with image_blocks as demo:
     warmup_status = gr.Textbox(visible=False)
     
     # 앱 로드 시 GPU Warm-up 자동 실행 (torch.compile 첫 컴파일)
-    demo.load(
-        fn=warmup_gpu,
-        inputs=None,
-        outputs=warmup_status,
-    )
+    if is_compile_for_zeroGPU == True:
+        print("✓ GPU warm-up is disabled for ZeroGPU")
+    else:
+        demo.load(
+            fn=warmup_gpu,
+            inputs=None,
+            outputs=warmup_status,
+        )
     
     print("✓ Gradio interface components created")
     print("✓ Event handlers configured")