Update app.py

app.py

@@ -22,14 +22,14 @@ os.makedirs(TMP_DIR, exist_ok=True)
 
 
 # 메모리 관련 환경 변수
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'  # 더 작은 값으로 설정
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['TORCH_HOME'] = '/tmp/torch_home'
 os.environ['HF_HOME'] = '/tmp/huggingface'
 os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
+os.environ['SPCONV_ALGO'] = 'native'
+os.environ['WARP_USE_CPU'] = '1'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
@@ -40,25 +40,17 @@ def initialize_models():
         # 메모리 설정
         torch.backends.cudnn.benchmark = False
         torch.backends.cudnn.deterministic = True
-        torch.backends.cuda.matmul.allow_tf32 = False
-        torch.backends.cudnn.allow_tf32 = False
-        
-        # 캐시 디렉토리 생성 및 정리
-        for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
-            os.makedirs(dir_path, exist_ok=True)
-            for file in os.listdir(dir_path):
-                try:
-                    os.remove(os.path.join(dir_path, file))
-                except:
-                    pass
         
+        print("Initializing Trellis pipeline...")
         # Trellis 파이프라인 초기화
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            device_map="auto",
-            torch_dtype=torch.float16  # 반정밀도 사용
+            "JeffreyXiang/TRELLIS-image-large"
         )
         
+        if pipeline is None:
+            raise Exception("Failed to initialize Trellis pipeline")
+            
+        print("Initializing translator...")
         # 번역기 초기화
         translator = translation_pipeline(
             "translation",
@@ -66,14 +58,17 @@ def initialize_models():
             device="cpu"
         )
         
+        if translator is None:
+            raise Exception("Failed to initialize translator")
+        
+        # Flux 파이프라인은 나중에 초기화
         flux_pipe = None
-        free_memory()
+        
         print("Models initialized successfully")
         return True
         
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
-        free_memory()
         return False
 
 def get_flux_pipe():
@@ -84,11 +79,8 @@ def get_flux_pipe():
             free_memory()
             flux_pipe = FluxPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-dev",
-                device_map="auto",
-                torch_dtype=torch.float16,
-                variant="fp16",
-                use_safetensors=True,
-                low_cpu_mem_usage=True
+                torch_dtype=torch.float32,  # CPU 모드로 시작
+                use_safetensors=True
             )
         except Exception as e:
             print(f"Error loading Flux pipeline: {e}")
@@ -142,6 +134,9 @@ def translate_if_korean(text):
 @spaces.GPU
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     try:
+        if pipeline is None:
+            raise Exception("Pipeline not initialized")
+            
         trial_id = str(uuid.uuid4())
         
         # 이미지가 너무 작은 경우 크기 조정
@@ -150,10 +145,18 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
             ratio = min_size / min(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
             
-        processed_image = pipeline.preprocess_image(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
         
     except Exception as e:
         print(f"Error in preprocess_image: {str(e)}")
@@ -439,21 +442,52 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
 
 if __name__ == "__main__":
     import warnings
-    warnings.filterwarnings('ignore')  # 경고 메시지 무시
+    warnings.filterwarnings('ignore')
+    
+    # 디렉토리 생성
+    os.makedirs(TMP_DIR, exist_ok=True)
     
+    # 메모리 정리
     free_memory()
     
-    if not initialize_models():
-        print("Failed to initialize models")
+    # 모델 초기화 시도
+    retry_count = 3
+    initialized = False
+    
+    for i in range(retry_count):
+        try:
+            if initialize_models():
+                initialized = True
+                break
+            else:
+                print(f"Initialization attempt {i+1} failed, retrying...")
+                free_memory()
+        except Exception as e:
+            print(f"Error during initialization attempt {i+1}: {str(e)}")
+            free_memory()
+    
+    if not initialized:
+        print("Failed to initialize models after multiple attempts")
         exit(1)
     
+    try:
+        # rembg 사전 로드 시도
+        test_image = Image.fromarray(np.ones((32, 32, 3), dtype=np.uint8) * 255)
+        if pipeline is not None:
+            pipeline.preprocess_image(test_image)
+    except Exception as e:
+        print(f"Warning: Failed to preload rembg: {str(e)}")
+    
     # Gradio 앱 실행
-    demo.queue(max_size=1).launch(  # 큐 크기를 1로 제한
-        share=True,
-        max_threads=1,
-        show_error=True,
-        enable_queue=True,
-        server_port=7860,
-        server_name="0.0.0.0",
-        quiet=True  # 로그 출력 최소화
-    )
+    try:
+        demo.queue(max_size=1).launch(
+            share=True,
+            max_threads=1,
+            show_error=True,
+            server_port=7860,
+            server_name="0.0.0.0",
+            quiet=True
+        )
+    except Exception as e:
+        print(f"Error launching Gradio app: {str(e)}")
+        exit(1)