Update app.py

app.py

@@ -25,7 +25,6 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import gc
 import torch.cuda.amp as amp
 
-# 메모리 관리 함수 강화
 def clear_memory():
     """메모리 정리 함수"""
     if torch.cuda.is_available():
@@ -34,7 +33,11 @@ def clear_memory():
     gc.collect()
 
 # 자동 혼합 정밀도(Automatic Mixed Precision) 설정
-scaler = amp.GradScaler()
+if torch.cuda.is_available():
+    scaler = torch.amp.GradScaler('cuda')
+else:
+    scaler = None
+
 
 
 model_name = "Helsinki-NLP/opus-mt-ko-en"
@@ -85,17 +88,16 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
+
 # FLUX 파이프라인 초기화
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     torch_dtype=torch.float16,  # A100에 최적화된 float16 사용
-    use_auth_token=HF_TOKEN,
-    device_map="balanced"
+    use_auth_token=HF_TOKEN
 )
 pipe.enable_attention_slicing(slice_size="auto")  # 메모리 사용량 최적화
-pipe.enable_sequential_cpu_offload()  # CPU 오프로딩 활성화
-
 
+# LoRA 가중치 로드
 pipe.load_lora_weights(
     hf_hub_download(
         "ByteDance/Hyper-SD",
@@ -105,6 +107,13 @@ pipe.load_lora_weights(
 )
 pipe.fuse_lora(lora_scale=0.125)
 
+# GPU 메모리 최적화
+if torch.cuda.is_available():
+    pipe.to("cuda")
+    pipe.enable_vae_slicing()  # VAE 슬라이싱 활성화
+    
+
+
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 단일 GPU 사용
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"  # CUDA 메모리 할당 설정
 
@@ -195,17 +204,19 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             width, height = adjust_size_to_multiple_of_8(width, height)
 
         with timer("Background generation"):
-            with torch.cuda.amp.autocast():  # 자동 혼합 정밀도 사용
-                image = pipe(
-                    prompt=prompt,
-                    width=width,
-                    height=height,
-                    num_inference_steps=8,
-                    guidance_scale=4.0,
-                    max_length=77,
-                ).images[0]
-                
-        return image
+            with torch.inference_mode():  # inference_mode 사용
+                with torch.cuda.amp.autocast():
+                    image = pipe(
+                        prompt=prompt,
+                        width=width,
+                        height=height,
+                        num_inference_steps=8,
+                        guidance_scale=4.0,
+                        max_length=77,
+                    ).images[0]
+                    
+            clear_memory()  # 즉시 메모리 정리
+            return image
     except Exception as e:
         print(f"Background generation error: {str(e)}")
         return Image.new('RGB', (512, 512), 'white')