Update app.py

app.py

@@ -36,38 +36,38 @@ g = GlobalVars()
 def initialize_models(device):
     try:
         print("Initializing models...")
-        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")
+        # 3D 생성 파이프라인
+        g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
+            "JeffreyXiang/TRELLIS-image-large",
+            torch_dtype=torch.float32  # 명시적으로 dtype 지정
+        )
+        
+        # 이미지 생성 파이프라인
+        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)}")
@@ -78,6 +78,7 @@ 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'
@@ -85,6 +86,7 @@ os.environ['SPARSE_BACKEND'] = 'native'
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['TORCH_USE_CUDA_DSA'] = '1'
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'  # xformers 관련 설정 추가
 
 # CUDA 초기화 방지
 torch.set_grad_enabled(False)
@@ -215,36 +217,46 @@ 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}")
         
-        with torch.cuda.device(0):
-            outputs = g.trellis_pipeline.run(
-                image,
-                seed=seed,
-                formats=["gaussian", "mesh"],
-                preprocess_image=False,
-                sparse_structure_sampler_params={
-                    "steps": ss_sampling_steps,
-                    "cfg_strength": ss_guidance_strength,
-                },
-                slat_sampler_params={
-                    "steps": slat_sampling_steps,
-                    "cfg_strength": slat_guidance_strength,
-                },
-            )
+        # Move model to GPU
+        g.trellis_pipeline.to("cuda")
+        
+        try:
+            with torch.inference_mode():
+                with torch.cuda.amp.autocast():
+                    outputs = g.trellis_pipeline.run(
+                        image,
+                        seed=seed,
+                        formats=["gaussian", "mesh"],
+                        preprocess_image=False,
+                        sparse_structure_sampler_params={
+                            "steps": ss_sampling_steps,
+                            "cfg_strength": ss_guidance_strength,
+                        },
+                        slat_sampler_params={
+                            "steps": slat_sampling_steps,
+                            "cfg_strength": slat_guidance_strength,
+                        },
+                    )
 
-            video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-            video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-            video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-            new_trial_id = str(uuid.uuid4())
-            video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
-            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)
-            
+                    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+                    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+                    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+                    new_trial_id = str(uuid.uuid4())
+                    video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
+                    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)
+                    
+                    return state, video_path
+        finally:
+            # Move model back to CPU
+            g.trellis_pipeline.to("cpu")
             torch.cuda.empty_cache()
-            return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
+        if hasattr(g.trellis_pipeline, 'device') and g.trellis_pipeline.device.type == "cuda":
+            g.trellis_pipeline.to("cpu")
         torch.cuda.empty_cache()
         return None, None