Update app.py

app.py

@@ -52,16 +52,18 @@ def elapsed_time(fn, *args):
 
 def main():
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-
-    elapsed, pipeline = elapsed_time(load_pipeline_w_lora)
-    pipeline = pipeline.to(device)
+    
 
     st.write(f"Model is loaded in {elapsed} seconds!")
     
     st.title("BAYC Text to IMAGE generator")
     st.write(f"Stable diffusion model is fine-tuned by lora using dataset {DATASET_REPO}")
 
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    elapsed, pipeline = elapsed_time(load_pipeline_w_lora)
+    pipeline = pipeline.to(device)
+
     sample = Image.open(SAMPLE_IMAGE)
     st.image(sample, caption="Example image with prompt <An ape with solid gold fur and beanie>")