ginipick commited on
Commit
3cee212
·
verified ·
1 Parent(s): e6920c4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -79,11 +79,11 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
79
  gd_model = gd_model.to(device=device)
80
  assert isinstance(gd_model, GroundingDinoForObjectDetection)
81
 
82
- # FLUX 파이프라인 초기화
83
  pipe = FluxPipeline.from_pretrained(
84
  "black-forest-labs/FLUX.1-dev",
85
- torch_dtype=torch.bfloat16,
86
- use_auth_token=HF_TOKEN
 
87
  )
88
  pipe.load_lora_weights(
89
  hf_hub_download(
@@ -93,7 +93,7 @@ pipe.load_lora_weights(
93
  )
94
  )
95
  pipe.fuse_lora(lora_scale=0.125)
96
- pipe.to(device="cuda", dtype=torch.bfloat16)
97
 
98
  class timer:
99
  def __init__(self, method_name="timed process"):
 
79
  gd_model = gd_model.to(device=device)
80
  assert isinstance(gd_model, GroundingDinoForObjectDetection)
81
 
 
82
  pipe = FluxPipeline.from_pretrained(
83
  "black-forest-labs/FLUX.1-dev",
84
+ torch_dtype=torch.float16, # 메모리 사용량 감소를 위해 float16 사용
85
+ use_auth_token=HF_TOKEN,
86
+ device_map="auto" # 자동 장치 매핑
87
  )
88
  pipe.load_lora_weights(
89
  hf_hub_download(
 
93
  )
94
  )
95
  pipe.fuse_lora(lora_scale=0.125)
96
+
97
 
98
  class timer:
99
  def __init__(self, method_name="timed process"):