Update app.py

app.py

@@ -23,6 +23,7 @@ from transformers import (
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Alpha-VLLM/Lumina-Image-2.0"
+transformer_repo_id = "benjamin-paine/Lumina-Image-2.0" # Temporarily fixed, change when main repo gets updated
 
 if torch.cuda.is_available():
     torch_dtype = torch.bfloat16
@@ -30,10 +31,9 @@ else:
     torch_dtype = torch.float32
 
 ###
-
+transformer = Lumina2Transformer2DModel.from_pretrained(transformer_repo_id, subfolder="transformer")
 vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
 text_encoder = Gemma2Model.from_pretrained(model_repo_id, subfolder="text_encoder")
-transformer = Lumina2Transformer2DModel.from_pretrained(model_repo_id, subfolder="transformer")
 tokenizer = GemmaTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler")