Add OUTPUT_DIR env var to inference script"

Author: bhavya01

examples/research_projects/pytorch_xla/training/text_to_image/inference_sdxl.py

@@ -9,12 +9,13 @@
 import torch_xla.runtime as xr
 
 CACHE_DIR = os.environ.get("CACHE_DIR", '/mnt/bbahl/xla_cache/')
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", '/mnt/bbahl/trained-model/')
 if CACHE_DIR:
     xr.initialize_cache(CACHE_DIR, readonly=False)
 
 
 device = xm.xla_device()
-model_path = "/mnt/bbahl/trained-model"
+model_path = OUTPUT_DIR
 pipe = StableDiffusionXLPipeline.from_pretrained(
     model_path, 
     torch_dtype=torch.bfloat16