Update app.py

app.py

@@ -21,10 +21,12 @@ TMP_DIR = "/tmp/Trellis-demo"
 os.makedirs(TMP_DIR, exist_ok=True)
 
 
-# GPU 메모리 관련 환경 변수 수정
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'  # A100에 맞게 증가
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # 단일 GPU 사용
-os.environ['CUDA_LAUNCH_BLOCKING'] = '0'  # A100에서는 비동기 실행 허용
+# 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['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
+os.environ['CUDA_CACHE_DISABLE'] = '1'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
@@ -44,15 +46,13 @@ def initialize_models():
         
         if torch.cuda.is_available():
             pipeline = pipeline.to("cuda")
-            # 모델을 FP16으로 변환
-            for param in pipeline.parameters():
-                param.data = param.data.half()
+            # FP16 변환은 제거 (pipeline이 자체적으로 처리)
             
         print("Initializing translator...")
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
-            device="cuda"
+            device="cuda" if torch.cuda.is_available() else "cpu"
         )
         
         # Flux 파이프라인은 나중에 초기화
@@ -74,9 +74,10 @@ def get_flux_pipe():
             flux_pipe = FluxPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-dev",
                 use_safetensors=True
-            ).to("cuda")
-            # FP16으로 변환
-            flux_pipe.to(torch.float16)
+            )
+            if torch.cuda.is_available():
+                flux_pipe = flux_pipe.to("cuda")
+                flux_pipe.enable_model_cpu_offload()  # CPU 오프로딩 활성화
         except Exception as e:
             print(f"Error loading Flux pipeline: {e}")
             return None
@@ -204,7 +205,7 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         input_image = Image.open(f"{TMP_DIR}/{trial_id}.png")
         
         # L40S에 맞게 이미지 크기 제한 조정
-        max_size = 768  # L40S는 더 큰 이미지 처리 가능
+        max_size = 768
         if max(input_image.size) > max_size:
             ratio = max_size / max(input_image.size)
             input_image = input_image.resize(
@@ -216,14 +217,14 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         if torch.cuda.is_available():
             pipeline.to("cuda")
         
-        with torch.cuda.amp.autocast():  # 자동 혼합 정밀도 사용
+        try:
             outputs = pipeline.run(
                 input_image,
                 seed=seed,
                 formats=["gaussian", "mesh"],
                 preprocess_image=False,
                 sparse_structure_sampler_params={
-                    "steps": min(ss_sampling_steps, 20),  # L40S에서 더 많은 스텝 허용
+                    "steps": min(ss_sampling_steps, 20),
                     "cfg_strength": ss_guidance_strength,
                 },
                 slat_sampler_params={
@@ -231,6 +232,10 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
                     "cfg_strength": slat_guidance_strength,
                 }
             )
+        except RuntimeError as e:
+            print(f"Runtime error in pipeline.run: {str(e)}")
+            free_memory()
+            raise e
         
         # 비디오 생성
         video = render_utils.render_video(outputs['gaussian'][0], num_frames=40)['color']