Update app.py

app.py

@@ -21,7 +21,16 @@ from diffusers import FluxPipeline
 from transformers import pipeline
 from typing import Tuple, Dict, Any  # Tuple import 추가
 
+# CUDA 메모리 관리 설정
+torch.cuda.empty_cache()
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
 
+# 환경 변수 설정
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+os.environ['SPCONV_ALGO'] = 'native'
+os.environ['SPARSE_BACKEND'] = 'native'
 
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -328,15 +337,15 @@ if __name__ == "__main__":
         pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.bfloat16,
-            use_auth_token=HF_TOKEN,  # token 대신 use_auth_token 사용
-            device_map="auto"
+            use_auth_token=HF_TOKEN,
+            device_map="balanced"  # "auto"에서 "balanced"로 변경
         )
         
         # Hyper-SD LoRA 로드
         lora_path = hf_hub_download(
             "ByteDance/Hyper-SD",
             "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-            use_auth_token=HF_TOKEN  # token 대신 use_auth_token 사용
+            use_auth_token=HF_TOKEN
         )
         pipe.load_lora_weights(lora_path)
         pipe.fuse_lora(lora_scale=0.125)