Update app.py

app.py

@@ -20,42 +20,39 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = "/tmp/Trellis-demo"
 os.makedirs(TMP_DIR, exist_ok=True)
 
-
-# GPU 메모리 관련 환경 변수
+# GPU 메모리 관련 환경 변수 추가
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # 디버깅을 위해 1로 설정
 os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
-os.environ['CUDA_CACHE_DISABLE'] = '1'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
     
     try:
-        import torch
-        
-        # L40S GPU 최적화 설정
-        torch.backends.cudnn.benchmark = True
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.backends.cudnn.allow_tf32 = True
+        # CUDA 설정
+        if torch.cuda.is_available():
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
         
         print("Initializing Trellis pipeline...")
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
+            "JeffreyXiang/TRELLIS-image-large",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
         )
         
         if torch.cuda.is_available():
             pipeline = pipeline.to("cuda")
-            # FP16 변환은 제거 (pipeline이 자체적으로 처리)
-            
+            pipeline.enable_model_cpu_offload()
+        
         print("Initializing translator...")
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
-            device="cuda" if torch.cuda.is_available() else "cpu"
+            device=0 if torch.cuda.is_available() else -1
         )
         
-        # Flux 파이프라인은 나중에 초기화
         flux_pipe = None
         
         print("Models initialized successfully")
@@ -128,30 +125,29 @@ def translate_if_korean(text):
 
 
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
+    if image is None:
+        return None, None
+        
     try:
-        if pipeline is None:
-            raise Exception("Pipeline not initialized")
-            
         trial_id = str(uuid.uuid4())
         
-        # 이미지가 너무 작은 경우 크기 조정
-        min_size = 64
-        if image.size[0] < min_size or image.size[1] < min_size:
-            ratio = min_size / min(image.size)
+        # 이미지 크기 제한
+        max_size = 768
+        if max(image.size) > max_size:
+            ratio = max_size / max(image.size)
             new_size = tuple(int(dim * ratio) for dim in image.size)
             image = image.resize(new_size, Image.LANCZOS)
-        
-        try:
-            processed_image = pipeline.preprocess_image(image)
-            if processed_image is None:
-                raise Exception("Failed to process image")
-                
-            processed_image.save(f"{TMP_DIR}/{trial_id}.png")
-            return trial_id, processed_image
             
-        except Exception as e:
-            print(f"Error in image preprocessing: {str(e)}")
-            return None, None
+        # 이미지 전처리
+        processed_image = pipeline.preprocess_image(image)
+        if processed_image is None:
+            raise Exception("Failed to process image")
+            
+        # 임시 파일 저장
+        save_path = os.path.join(TMP_DIR, f"{trial_id}.png")
+        processed_image.save(save_path)
+        
+        return trial_id, processed_image
         
     except Exception as e:
         print(f"Error in preprocess_image: {str(e)}")
@@ -452,15 +448,21 @@ if __name__ == "__main__":
     # 메모리 정리
     free_memory()
     
-    # 모델 초기화
-    if not initialize_models():
-        print("Failed to initialize models")
+    # 모델 초기화 - 여러 번 시도
+    max_attempts = 3
+    for attempt in range(max_attempts):
+        if initialize_models():
+            break
+        print(f"Attempt {attempt + 1} failed, retrying...")
+        free_memory()
+    else:
+        print("Failed to initialize models after multiple attempts")
         exit(1)
     
     # Gradio 앱 실행
-    demo.queue(max_size=2).launch(  # 큐 크기 증가
+    demo.queue(max_size=2).launch(
         share=True,
-        max_threads=4,  # 스레드 수 증가
+        max_threads=4,
         show_error=True,
         server_port=7860,
         server_name="0.0.0.0"