Update app.py

app.py

@@ -18,11 +18,10 @@ from typing import *
 
 # 메모리 관련 환경 변수
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
 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'
@@ -36,23 +35,26 @@ def initialize_models():
     global pipeline, translator, flux_pipe
     
     try:
-        # Trellis 파이프라인 초기화 (기본 CPU 모드)
+        # 캐시 디렉토리 생성
+        for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
+            os.makedirs(dir_path, exist_ok=True)
+        
+        # Trellis 파이프라인 초기화 (메모리 최적화)
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
+            "JeffreyXiang/TRELLIS-image-large",
+            resume_download=True,
+            local_files_only=False
         )
         
-        # 번역기 초기화
+        # 번역기 초기화 (더 작은 모델 사용)
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
             device="cpu"
         )
         
-        # Flux 파이프라인 초기화
-        flux_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.float32  # CPU 모드
-        )
+        # Flux 파이프라인 초기화 (메모리 최적화)
+        flux_pipe = None  # 필요할 때 로드
         
         print("Models initialized successfully")
         return True
@@ -61,10 +63,31 @@ def initialize_models():
         print(f"Model initialization error: {str(e)}")
         return False
 
+def load_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
+        )
+    return flux_pipe
+
 def free_memory():
     """메모리를 정리하는 유틸리티 함수"""
     import gc
     gc.collect()
+    
+    # 임시 파일 정리
+    for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
+        if os.path.exists(dir_path):
+            for file in os.listdir(dir_path):
+                file_path = os.path.join(dir_path, file)
+                try:
+                    if os.path.isfile(file_path):
+                        os.unlink(file_path)
+                except Exception as e:
+                    print(f'Error deleting {file_path}: {e}')
 
 @spaces.GPU
 def setup_gpu_model(model):
@@ -73,29 +96,6 @@ def setup_gpu_model(model):
         model = model.to("cuda")
     return model
 
-if __name__ == "__main__":
-    # CPU 메모리만 정리
-    free_memory()
-    
-    # 모델 초기화
-    if not initialize_models():
-        print("Failed to initialize models")
-        exit(1)
-    
-    try:
-        # rembg 사전 로드 시도 (작은 이미지로)
-        test_image = Image.fromarray(np.ones((64, 64, 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=5).launch(
-        share=True,
-        max_threads=2,
-        show_error=True,
-        cache_examples=False
-    )
 
 def translate_if_korean(text):
     if any(ord('가') <= ord(char) <= ord('힣') for char in text):
@@ -225,18 +225,18 @@ 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:
-        # GPU 설정
+        # Flux 파이프라인 필요할 때 로드
+        flux_pipe = load_flux_pipe()
+        
         if torch.cuda.is_available():
             flux_pipe.to("cuda")
-            flux_pipe.to(torch.float16)
         
-        # 기본 프롬프트를 추가
-        base_prompt = "wbgmsst, 3D, white background"
+        # 이미지 크기 제한
+        height = min(height, 512)
+        width = min(width, 512)
         
-        # 사용자 프롬프트를 번역 (한국어인 경우)
+        base_prompt = "wbgmsst, 3D, white background"
         translated_prompt = translate_if_korean(prompt)
-        
-        # 최종 프롬프트 조합
         final_prompt = f"{translated_prompt}, {base_prompt}"
         
         with torch.inference_mode():
@@ -244,18 +244,19 @@ def generate_image_from_text(prompt, height, width, guidance_scale, num_steps):
                 prompt=[final_prompt],
                 height=height,
                 width=width,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_steps
+                guidance_scale=min(guidance_scale, 10.0),
+                num_inference_steps=min(num_steps, 30)
             ).images[0]
         
-        # CPU 모드로 돌아가기
-        flux_pipe.to("cpu")
+        if torch.cuda.is_available():
+            flux_pipe.to("cpu")
         
         return image
         
     except Exception as e:
         print(f"Error in generate_image_from_text: {str(e)}")
-        flux_pipe.to("cpu")
+        if torch.cuda.is_available() and flux_pipe is not None:
+            flux_pipe.to("cpu")
         raise e
 
 @spaces.GPU
@@ -395,7 +396,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     )
 
 if __name__ == "__main__":
-    # CPU 메모리만 정리
+    # 메모리 정리
     free_memory()
     
     # 모델 초기화
@@ -404,16 +405,16 @@ if __name__ == "__main__":
         exit(1)
     
     try:
-        # rembg 사전 로드 시도 (작은 이미지로)
-        test_image = Image.fromarray(np.ones((64, 64, 3), dtype=np.uint8) * 255)
+        # 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=5).launch(
+    demo.queue(max_size=3).launch(
         share=True,
-        max_threads=2,
+        max_threads=1,
         show_error=True,
         cache_examples=False,
         enable_queue=True