load model in session

app.py

@@ -4,7 +4,7 @@ import logging
 import gradio as gr
 logging.basicConfig(level=logging.INFO)
 from src.utils import generate_centered_gaussian_noise
-from src.demo import resize,plot_flow,plot_diff,load_model_diff,load_model_flow_localized,load_model_flow_standard
+from src.demo import resize,plot_flow,plot_diff,load_models,load_model_diff,load_model_flow_localized,load_model_flow_standard
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 img_shape = (1, 28, 28)
@@ -17,15 +17,15 @@ alphas = 1.0 - betas
 alphas_cumprod = torch.cumprod(alphas, dim=0).to(device)
 
 
-#model_diff,model_flow_standard,model_flow_localized = load_models(ENV,device=device)
-#not catching models because of memory limit in free deployment
+model_diff,model_flow_standard,model_flow_localized = load_models(ENV,device=device)
+###not catching models because of memory limit in free deployment
 
 @torch.no_grad()
 def generate_diffusion_intermediates_streaming(label):
     logging.info("🚀 Starting Diffusion Generation")
     total_start = time.time()
 
-    model_diff = load_model_diff(ENV,device=device)
+    #model_diff = load_model_diff(ENV,device=device)
 
 
     x = torch.randn(1, *img_shape).to(device)
@@ -90,11 +90,11 @@ def generate_flow_intermediates_streaming(label, noise_type):
     # Select noise and model
     if noise_type == "Localized":
         x = generate_centered_gaussian_noise((1, *img_shape)).to(device)
-        model_flow =  load_model_flow_localized(ENV,device=device)
+        model_flow = model_flow_localized # load_model_flow_localized(ENV,device=device)
 
     else:
         x = torch.randn(1, *img_shape).to(device)
-        model_flow =  load_model_flow_standard(ENV,device=device)
+        model_flow =  model_flow_standard #load_model_flow_standard(ENV,device=device)
 
     y = torch.full((1,), label, dtype=torch.long, device=device)
     steps = 50