Update app.py

app.py

@@ -77,7 +77,7 @@
 
 # if __name__ == "__main__":
 #     demo.launch()
-########################3rd######################3
+########################3rd-MAIN######################3
 
 # import torch
 # import gradio as gr
@@ -138,7 +138,65 @@
 
 
 #################4th##################
+
+# import torch
+# import gradio as gr
+# import requests
+# import os
+
+# # Define model repo
+# model_repo = "Mariam-Elz/CRM"
+
+# # Define model files and download paths
+# model_files = {
+#     "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():
+#     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(output_path, "wb") as f:
+#             f.write(response.content)
+
+# # Load model with low memory usage
+# def load_model():
+#     model_path = "models/CRM.pth"
+#     model = torch.load(model_path, map_location="cpu")  # Load on CPU to reduce memory usage
+#     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():
+#         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.squeeze(0).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()
+
+
+##############5TH#################
 import torch
+import torch.nn as nn
 import gradio as gr
 import requests
 import os
@@ -162,11 +220,26 @@ for filename, output_path in model_files.items():
         with open(output_path, "wb") as f:
             f.write(response.content)
 
-# Load model with low memory usage
+# Define the model architecture (you MUST replace this with your actual model)
+class CRM_Model(nn.Module):
+    def __init__(self):
+        super(CRM_Model, self).__init__()
+        self.layer1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.relu = nn.ReLU()
+        self.layer2 = nn.Conv2d(64, 3, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.relu(x)
+        x = self.layer2(x)
+        return x
+
+# Load model with proper architecture
 def load_model():
+    model = CRM_Model()  # Instantiate the model architecture
     model_path = "models/CRM.pth"
-    model = torch.load(model_path, map_location="cpu")  # Load on CPU to reduce memory usage
-    model.eval()
+    model.load_state_dict(torch.load(model_path, map_location="cpu"))  # Load weights
+    model.eval()  # Set to evaluation mode
     return model
 
 model = load_model()
@@ -175,10 +248,10 @@ model = load_model()
 def infer(image):
     """Process input image and return a reconstructed image."""
     with torch.no_grad():
-        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.squeeze(0).numpy()
+        image_tensor = torch.tensor(image).unsqueeze(0).permute(0, 3, 1, 2).float() / 255.0  # Convert to tensor
+        output = model(image_tensor)  # Run through model
+        output = output.squeeze(0).permute(1, 2, 0).numpy() * 255.0  # Convert back to image
+        return output.astype("uint8")
 
 # Create Gradio UI
 demo = gr.Interface(