Update app.py

app.py

@@ -79,37 +79,93 @@
 #     demo.launch()
 ########################3rd######################3
 
+# import torch
+# import gradio as gr
+# import requests
+# import os
+
+# # Download model weights from Hugging Face model repo (if not already present)
+# model_repo = "Mariam-Elz/CRM"  # Your Hugging Face model repo
+
+# model_files = {
+#     "ccm-diffusion.pth": "ccm-diffusion.pth",
+#     "pixel-diffusion.pth": "pixel-diffusion.pth",
+#     "CRM.pth": "CRM.pth",
+# }
+
+# os.makedirs("models", exist_ok=True)
+
+# for filename, output_path in model_files.items():
+#     file_path = f"models/{output_path}"
+#     if not os.path.exists(file_path):
+#         url = f"https://huggingface.co/{model_repo}/resolve/main/{filename}"
+#         print(f"Downloading {filename}...")
+#         response = requests.get(url)
+#         with open(file_path, "wb") as f:
+#             f.write(response.content)
+
+# # Load model (This part depends on how the model is defined)
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# def load_model():
+#     model_path = "models/CRM.pth"
+#     model = torch.load(model_path, map_location=device)
+#     model.eval()
+#     return model
+
+# model = load_model()
+
+# # Define inference function
+# def infer(image):
+#     """Process input image and return a reconstructed image."""
+#     with torch.no_grad():
+#         # Assuming model expects a tensor input
+#         image_tensor = torch.tensor(image).to(device)
+#         output = model(image_tensor)
+#         return output.cpu().numpy()
+
+# # Create Gradio UI
+# demo = gr.Interface(
+#     fn=infer,
+#     inputs=gr.Image(type="numpy"),
+#     outputs=gr.Image(type="numpy"),
+#     title="Convolutional Reconstruction Model",
+#     description="Upload an image to get the reconstructed output."
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+#################4th##################
 import torch
 import gradio as gr
 import requests
 import os
 
-# Download model weights from Hugging Face model repo (if not already present)
-model_repo = "Mariam-Elz/CRM"  # Your Hugging Face model repo
+# Define model repo
+model_repo = "Mariam-Elz/CRM"
 
+# Define model files and download paths
 model_files = {
-    "ccm-diffusion.pth": "ccm-diffusion.pth",
-    "pixel-diffusion.pth": "pixel-diffusion.pth",
-    "CRM.pth": "CRM.pth",
+    "CRM.pth": "models/CRM.pth"
 }
 
 os.makedirs("models", exist_ok=True)
 
+# Download model files only if they don't exist
 for filename, output_path in model_files.items():
-    file_path = f"models/{output_path}"
-    if not os.path.exists(file_path):
+    if not os.path.exists(output_path):
         url = f"https://huggingface.co/{model_repo}/resolve/main/{filename}"
         print(f"Downloading {filename}...")
         response = requests.get(url)
-        with open(file_path, "wb") as f:
+        with open(output_path, "wb") as f:
             f.write(response.content)
 
-# Load model (This part depends on how the model is defined)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
+# Load model with low memory usage
 def load_model():
     model_path = "models/CRM.pth"
-    model = torch.load(model_path, map_location=device)
+    model = torch.load(model_path, map_location="cpu")  # Load on CPU to reduce memory usage
     model.eval()
     return model
 
@@ -119,10 +175,10 @@ model = load_model()
 def infer(image):
     """Process input image and return a reconstructed image."""
     with torch.no_grad():
-        # Assuming model expects a tensor input
-        image_tensor = torch.tensor(image).to(device)
+        image_tensor = torch.tensor(image).unsqueeze(0)  # Add batch dimension
+        image_tensor = image_tensor.to("cpu")  # Keep on CPU to save memory
         output = model(image_tensor)
-        return output.cpu().numpy()
+        return output.squeeze(0).numpy()
 
 # Create Gradio UI
 demo = gr.Interface(