Update app.py

app.py

@@ -36,37 +36,38 @@ g = GlobalVars()
 def initialize_models(device):
     try:
         print("Initializing models...")
-        # 3D 생성 파이프라인
-        g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
-        )
-        
-        # 이미지 생성 파이프라인
-        print("Loading flux_pipe...")
-        g.flux_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.bfloat16,
-            device_map="balanced"
-        )
-        
-        # Hyper-SD LoRA 로드
-        print("Loading LoRA weights...")
-        lora_path = hf_hub_download(
-            "ByteDance/Hyper-SD",
-            "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-            use_auth_token=HF_TOKEN
-        )
-        g.flux_pipe.load_lora_weights(lora_path)
-        g.flux_pipe.fuse_lora(lora_scale=0.125)
-        
-        # 번역기 초기화
-        print("Initializing translator...")
-        g.translator = transformers_pipeline(
-            "translation", 
-            model="Helsinki-NLP/opus-mt-ko-en",
-            device=device
-        )
-        print("Model initialization completed successfully")
+        with torch.no_grad():
+            # 3D 생성 파이프라인
+            g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
+                "JeffreyXiang/TRELLIS-image-large"
+            )
+            
+            # 이미지 생성 파이프라인
+            print("Loading flux_pipe...")
+            g.flux_pipe = FluxPipeline.from_pretrained(
+                "black-forest-labs/FLUX.1-dev",
+                torch_dtype=torch.bfloat16,
+                device_map="balanced"
+            )
+            
+            # Hyper-SD LoRA 로드
+            print("Loading LoRA weights...")
+            lora_path = hf_hub_download(
+                "ByteDance/Hyper-SD",
+                "Hyper-FLUX.1-dev-8steps-lora.safetensors",
+                use_auth_token=HF_TOKEN
+            )
+            g.flux_pipe.load_lora_weights(lora_path)
+            g.flux_pipe.fuse_lora(lora_scale=0.125)
+            
+            # 번역기 초기화
+            print("Initializing translator...")
+            g.translator = transformers_pipeline(
+                "translation", 
+                model="Helsinki-NLP/opus-mt-ko-en",
+                device=device
+            )
+            print("Model initialization completed successfully")
         
     except Exception as e:
         print(f"Error during model initialization: {str(e)}")
@@ -77,13 +78,17 @@ 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'
 os.environ['TORCH_USE_CUDA_DSA'] = '1'
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+# CUDA 초기화 방지
+torch.set_grad_enabled(False)
+
 
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")