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,load_models,plot_diff
+from src.demo import resize,plot_flow,plot_diff,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)
@@ -16,14 +16,16 @@ betas = torch.linspace(1e-4, 0.02, timesteps)
 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)
 
+#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)
 
 
     x = torch.randn(1, *img_shape).to(device)
@@ -88,10 +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 = model_flow_localized
+        model_flow =  load_model_flow_localized(ENV,device=device)
+
     else:
         x = torch.randn(1, *img_shape).to(device)
-        model_flow = model_flow_standard
+        model_flow =  load_model_flow_standard(ENV,device=device)
 
     y = torch.full((1,), label, dtype=torch.long, device=device)
     steps = 50

src/demo.py

@@ -31,6 +31,42 @@ def load_models(ENV,device):
     model_flow_localized.eval()
 
     return model_diff_standard,model_flow_standard,model_flow_localized
+
+
+def load_model_diff(ENV,device):
+    if ENV=="DEPLOY":
+        model_path = hf_hub_download(repo_id="CristianLazoQuispe/MNIST_Diff_Flow_matching", filename="outputs/diffusion/diffusion_model.pth",cache_dir="models")
+    else:
+        model_path  = "outputs/diffusion/diffusion_model.pth"
+    print("Diff Downloaded!")
+    model_diff_standard  = ConditionalUNet().to(device)
+    model_diff_standard.load_state_dict(torch.load(model_path, map_location=device))
+    model_diff_standard.eval()
+    return model_diff_standard
+
+def load_model_flow_standard(ENV,device):
+    if ENV=="DEPLOY":
+        model_path_standard  = hf_hub_download(repo_id="CristianLazoQuispe/MNIST_Diff_Flow_matching", filename="outputs/flow_matching/flow_model.pth",cache_dir="models")
+    else:
+        model_path_standard  = "outputs/flow_matching/flow_model.pth"
+    print("Flow Downloaded!")
+    model_flow_standard  = ConditionalUNet().to(device)
+    model_flow_standard.load_state_dict(torch.load(model_path_standard, map_location=device))
+    model_flow_standard.eval()
+    return model_flow_standard
+
+def load_model_flow_localized(ENV,device):
+    if ENV=="DEPLOY":
+        model_path_localized = hf_hub_download(repo_id="CristianLazoQuispe/MNIST_Diff_Flow_matching", filename="outputs/flow_matching/flow_model_localized_noise.pth",cache_dir="models")
+    else:
+        model_path_localized = "outputs/flow_matching/flow_model_localized_noise.pth"
+    print("Flow Downloaded!")
+    model_flow_localized = ConditionalUNet().to(device)
+    model_flow_localized.load_state_dict(torch.load(model_path_localized, map_location=device))
+    model_flow_localized.eval()
+    return model_flow_localized
+
+
 def resize(image,size=(200,200)):
     stretch_near = cv2.resize(image, size, interpolation = cv2.INTER_LINEAR)
     return stretch_near