Update app.py

app.py

@@ -76,12 +76,14 @@ def initialize_models(device):
 torch.cuda.empty_cache()
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.benchmark = True
+
+
 # 환경 변수 설정
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
 os.environ['SPCONV_ALGO'] = 'native'
 os.environ['SPARSE_BACKEND'] = 'native'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # CUDA 작업 동기화
-
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+os.environ['TORCH_USE_CUDA_DSA'] = '1'
 
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -208,6 +210,9 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         image = Image.open(image_path)
         print(f"Successfully loaded image with size: {image.size}")
         
+        # GPU로 모델 이동
+        if not g.trellis_pipeline.device.type == "cuda":
+            g.trellis_pipeline = g.trellis_pipeline.to("cuda")
         
         outputs = g.trellis_pipeline.run(
             image,
@@ -232,10 +237,16 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         os.makedirs(os.path.dirname(video_path), exist_ok=True)
         imageio.mimsave(video_path, video, fps=15)
         state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
+        
+        # CPU로 모델 되돌리기
+        g.trellis_pipeline = g.trellis_pipeline.to("cpu")
+        
         return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
+        if g.trellis_pipeline.device.type == "cuda":
+            g.trellis_pipeline = g.trellis_pipeline.to("cpu")
         return None, None