Update app.py

app.py

@@ -118,16 +118,16 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
     dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.precision]
 
     text_encoder = AutoModel.from_pretrained(
-        gemma_path, torch_dtype=dtype, device_map="cuda", token=hf_token
+        gemma_path, torch_dtype=dtype, device_map="cuda", token=args.hf_token
     ).eval()
     cap_feat_dim = text_encoder.config.hidden_size
     if args.num_gpus > 1:
         raise NotImplementedError("Inference with >1 GPUs not yet supported")
 
-    tokenizer = AutoTokenizer.from_pretrained(gemma_path, token=hf_token)
+    tokenizer = AutoTokenizer.from_pretrained(gemma_path, token=args.hf_token)
     tokenizer.padding_side = "right"
 
-    vae = AutoencoderKL.from_pretrained(flux_path, subfolder="vae", token=hf_token).cuda()
+    vae = AutoencoderKL.from_pretrained(flux_path, subfolder="vae", token=args.hf_token).cuda()
 
     print(f"Creating DiT: {train_args.model}")