Cicici1109 commited on
Commit
70230ab
·
verified ·
1 Parent(s): 351b9c1

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +5 -2
utils.py CHANGED
@@ -34,9 +34,12 @@ model_dict = {}
34
  def init_flux_pipeline():
35
  global pipe
36
  if pipe is None:
 
 
 
37
  pipe = FluxPipeline.from_pretrained(
38
- "black-forest-labs/FLUX.1-schnell",
39
- use_auth_token=True,
40
  torch_dtype=torch.bfloat16
41
  )
42
  pipe = pipe.to("cuda")
 
34
  def init_flux_pipeline():
35
  global pipe
36
  if pipe is None:
37
+ token = os.getenv("HF_TOKEN")
38
+ if not token:
39
+ raise ValueError("HF_TOKEN environment variable not set.")
40
  pipe = FluxPipeline.from_pretrained(
41
+ "black-forest-labs/FLUX.1-dev",
42
+ use_auth_token=token,
43
  torch_dtype=torch.bfloat16
44
  )
45
  pipe = pipe.to("cuda")