Update app.py

app.py

@@ -51,7 +51,8 @@ from multiprocessing import Process,Queue,set_start_method,get_context
 class ModelFailure:
     pass
 
-gemma_path = "./gemma-2-2b"
+gemma_path = "/home/user/app/checkpoints/text_encoder"
+flux_path = "/home/user/app/checkpoints/vae"
 #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
 # Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt
 def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompts, is_train=True):
@@ -118,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(
-        "google/gemma-2-2b", torch_dtype=dtype, device_map="cuda", token=hf_token
+        gemma_path, 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("google/gemma-2-2b", token=hf_token)
+    tokenizer = AutoTokenizer.from_pretrained(gemma_path, token=hf_token)
     tokenizer.padding_side = "right"
 
-    vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", token=hf_token).cuda()
+    vae = AutoencoderKL.from_pretrained(flux_path, subfolder="vae", token=hf_token).cuda()
 
     print(f"Creating DiT: {train_args.model}")