Update app.py

app.py

@@ -25,8 +25,11 @@ default_system_prompt = "You are an assistant designed to generate superior imag
 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
-torch_dtype = torch.float32
-
+if torch.cuda.is_available():
+    torch_dtype = torch.bfloat16
+else:
+    torch_dtype = torch.float32
+    
 ###
 transformer = Lumina2Transformer2DModel.from_pretrained(transformer_repo_id, subfolder="transformer")
 vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
@@ -47,7 +50,7 @@ pipe.to(device, torch_dtype)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1536
 
-@spaces.GPU(duration=65)
+@spaces.GPU(duration=60)
 def infer(
     prompt,
     negative_prompt="",