Update app.py

app.py

@@ -20,26 +20,43 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = "/tmp/Trellis-demo"
 os.makedirs(TMP_DIR, exist_ok=True)
 
+
 # 메모리 관련 환경 변수
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
+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['SPCONV_ALGO'] = 'native'
-os.environ['WARP_USE_CPU'] = '1'
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
     
     try:
-        # 캐시 디렉토리 생성
+        import torch
+        
+        # 메모리 설정
+        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
         
         # Trellis 파이프라인 초기화
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
+            "JeffreyXiang/TRELLIS-image-large",
+            device_map="auto",
+            torch_dtype=torch.float16  # 반정밀도 사용
         )
         
         # 번역기 초기화
@@ -49,49 +66,64 @@ def initialize_models():
             device="cpu"
         )
         
-        # Flux 파이프라인 초기화 (즉시 로드)
-        flux_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.float32
-        )
-        
+        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 load_flux_pipe():
-    """Flux 파이프라인을 필요할 때만 로드"""
+def get_flux_pipe():
+    """Flux 파이프라인을 필요할 때만 로드하는 함수"""
     global flux_pipe
     if flux_pipe is None:
-        flux_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.float32
-        )
+        try:
+            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
+            )
+        except Exception as e:
+            print(f"Error loading Flux pipeline: {e}")
+            return None
     return flux_pipe
 
+
 def free_memory():
-    """메모리를 정리하는 유틸리티 함수"""
+    """강화된 메모리 정리 함수"""
     import gc
+    import os
+    
+    # Python 가비지 컬렉션
     gc.collect()
     
+    # CUDA 메모리 정리
     if torch.cuda.is_available():
-        with torch.cuda.device('cuda'):
-            torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
     
     # 임시 파일 정리
-    for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
+    tmp_dirs = ['/tmp/transformers_cache', '/tmp/torch_home', 
+                '/tmp/huggingface', '/tmp/cache']
+    
+    for dir_path in tmp_dirs:
         if os.path.exists(dir_path):
-            for file in os.listdir(dir_path):
-                file_path = os.path.join(dir_path, file)
-                try:
+            try:
+                for file in os.listdir(dir_path):
+                    file_path = os.path.join(dir_path, file)
                     if os.path.isfile(file_path):
-                        os.unlink(file_path)
-                except Exception as e:
-                    print(f'Error deleting {file_path}: {e}')
-
+                        try:
+                            os.unlink(file_path)
+                        except:
+                            pass
+            except:
+                pass
 
 @spaces.GPU
 def setup_gpu_model(model):
@@ -230,11 +262,12 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
 @spaces.GPU
 def generate_image_from_text(prompt, height, width, guidance_scale, num_steps):
     try:
-        global flux_pipe
+        free_memory()  # 시작 전 메모리 정리
         
-        if torch.cuda.is_available():
-            flux_pipe = flux_pipe.to("cuda")
-            flux_pipe = flux_pipe.to(torch.float16)
+        # Flux 파이프라인 가져오기
+        flux_pipe = get_flux_pipe()
+        if flux_pipe is None:
+            raise Exception("Failed to load Flux pipeline")
         
         # 이미지 크기 제한
         height = min(height, 512)
@@ -245,28 +278,22 @@ def generate_image_from_text(prompt, height, width, guidance_scale, num_steps):
         translated_prompt = translate_if_korean(prompt)
         final_prompt = f"{translated_prompt}, {base_prompt}"
         
-        print(f"Generating image with prompt: {final_prompt}")
-        
-        with torch.inference_mode():
+        with torch.inference_mode(), torch.cuda.amp.autocast():
             output = flux_pipe(
                 prompt=[final_prompt],
                 height=height,
                 width=width,
-                guidance_scale=min(guidance_scale, 10.0),
-                num_inference_steps=min(num_steps, 30)
+                guidance_scale=min(guidance_scale, 7.5),  # 낮은 값으로 제한
+                num_inference_steps=min(num_steps, 20)    # 스텝 수 제한
             )
-            
             image = output.images[0]
         
-        if torch.cuda.is_available():
-            flux_pipe = flux_pipe.to("cpu")
-        
+        free_memory()  # 완료 후 메모리 정리
         return image
         
     except Exception as e:
         print(f"Error in generate_image_from_text: {str(e)}")
-        if torch.cuda.is_available():
-            flux_pipe = flux_pipe.to("cpu")
+        free_memory()
         raise e
 
 @spaces.GPU
@@ -411,24 +438,22 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
 
 
 if __name__ == "__main__":
-    # 메모리 정리
+    import warnings
+    warnings.filterwarnings('ignore')  # 경고 메시지 무시
+    
     free_memory()
     
-    # 모델 초기화
     if not initialize_models():
         print("Failed to initialize models")
         exit(1)
     
-    try:
-        # rembg 사전 로드 시도 (매우 작은 이미지로)
-        test_image = Image.fromarray(np.ones((32, 32, 3), dtype=np.uint8) * 255)
-        pipeline.preprocess_image(test_image)
-    except Exception as e:
-        print(f"Warning: Failed to preload rembg: {str(e)}")
-    
     # Gradio 앱 실행
-    demo.queue(max_size=3).launch(
+    demo.queue(max_size=1).launch(  # 큐 크기를 1로 제한
         share=True,
         max_threads=1,
-        show_error=True
+        show_error=True,
+        enable_queue=True,
+        server_port=7860,
+        server_name="0.0.0.0",
+        quiet=True  # 로그 출력 최소화
     )