Update app.py

app.py

@@ -79,37 +79,94 @@
 #     demo.launch()
 ########################3rd-MAIN######################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 +176,10 @@
 # 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(
@@ -137,9 +194,9 @@
 #     demo.launch()
 
 
-#################4th##################
-
+# ##############5TH#################
 # import torch
+# import torch.nn as nn
 # import gradio as gr
 # import requests
 # import os
@@ -163,11 +220,26 @@
 #         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()
@@ -176,10 +248,10 @@
 # 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(
@@ -192,75 +264,3 @@
 
 # if __name__ == "__main__":
 #     demo.launch()
-
-
-##############5TH#################
-import torch
-import torch.nn as nn
-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)
-
-# 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.load_state_dict(torch.load(model_path, map_location="cpu"))  # Load weights
-    model.eval()  # Set to evaluation mode
-    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).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(
-    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()