Update app.py

app.py

@@ -18,7 +18,7 @@ snapshot_download(
     repo_id="Alpha-VLLM/Lumina-Image-2.0", local_dir="/home/user/app/checkpoints"
 )
 
-
+hf_token = os.environ["HF_TOKEN"]
 
 import argparse
 import os
@@ -119,16 +119,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=args.hf_token
+        "google/gemma-2-2b", torch_dtype=dtype, device_map="cuda", token=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=args.hf_token)
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b", token=hf_token)
     tokenizer.padding_side = "right"
 
-    vae = AutoencoderKL.from_pretrained(flux_path, subfolder="vae", token=args.hf_token).cuda()
+    vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", token=hf_token).cuda()
 
     print(f"Creating DiT: {train_args.model}")