Test

app.py

@@ -9,10 +9,10 @@ import streamlit as st
 
 # model_path = WEIGHTS_DIR             # If you want to use previously trained model saved in gdrive, replace this with the full path of model in gdrive
 
-# pipe = StableDiffusionPipeline.from_pretrained(model_path, safety_checker=None, torch_dtype=torch.float32).to("cuda")
-# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-# pipe.enable_xformers_memory_efficient_attention()
-# g_cuda = None
+pipe = StableDiffusionPipeline.from_pretrained("Divyanshu04/Finetuned-sd-vae", safety_checker=None, torch_dtype=torch.float32).to("cuda")
+pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe.enable_xformers_memory_efficient_attention()
+g_cuda = None
          
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory