Update app.py

app.py

@@ -79,11 +79,11 @@ 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.bfloat16,
-    use_auth_token=HF_TOKEN
+    torch_dtype=torch.float16,  # 메모리 사용량 감소를 위해 float16 사용
+    use_auth_token=HF_TOKEN,
+    device_map="auto"  # 자동 장치 매핑
 )
 pipe.load_lora_weights(
     hf_hub_download(
@@ -93,7 +93,7 @@ pipe.load_lora_weights(
     )
 )
 pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device="cuda", dtype=torch.bfloat16)
+
 
 class timer:
     def __init__(self, method_name="timed process"):