Update app.py

app.py

@@ -79,62 +79,62 @@
 #     demo.launch()
 ########################3rd-MAIN######################3
 
-import torch
-import gradio as gr
-import requests
-import os
+# 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
+# # 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",
-}
+# model_files = {
+#     "ccm-diffusion.pth": "ccm-diffusion.pth",
+#     "pixel-diffusion.pth": "pixel-diffusion.pth",
+#     "CRM.pth": "CRM.pth",
+# }
 
-os.makedirs("models", exist_ok=True)
+# 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)
+# 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"
+# # 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
+# def load_model():
+#     model_path = "models/CRM.pth"
+#     model = torch.load(model_path, map_location=device)
+#     model.eval()
+#     return model
 
-model = load_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()
+# # 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."
-)
+# # 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()
+# if __name__ == "__main__":
+#     demo.launch()
 
 
 #################4th##################
@@ -264,3 +264,74 @@ if __name__ == "__main__":
 
 # if __name__ == "__main__":
 #     demo.launch()
+
+
+#############6th##################
+import torch
+import gradio as gr
+import requests
+import os
+import numpy as np
+
+# Hugging Face Model Repository
+model_repo = "Mariam-Elz/CRM"
+
+# Download Model Weights (Only CRM.pth to Save Memory)
+model_path = "models/CRM.pth"
+os.makedirs("models", exist_ok=True)
+
+if not os.path.exists(model_path):
+    url = f"https://huggingface.co/{model_repo}/resolve/main/CRM.pth"
+    print(f"Downloading CRM.pth...")
+    response = requests.get(url)
+    with open(model_path, "wb") as f:
+        f.write(response.content)
+
+# Set Device (Use CPU to Reduce RAM Usage)
+device = "cpu"
+
+# Load Model Efficiently
+def load_model():
+    model = torch.load(model_path, map_location=device)
+    if isinstance(model, torch.nn.Module):
+        model.eval()  # Ensure model is in inference mode
+    return model
+
+# Load model only when needed (saves memory)
+model = load_model()
+
+# Define Inference Function with Memory Optimizations
+def infer(image):
+    """Process input image and return a reconstructed image."""
+    with torch.no_grad():
+        # Convert image to torch tensor & normalize (float16 to save RAM)
+        image_tensor = torch.tensor(image, dtype=torch.float16).unsqueeze(0).permute(0, 3, 1, 2) / 255.0
+        image_tensor = image_tensor.to(device)
+
+        # Model Inference
+        output = model(image_tensor)
+
+        # Convert back to numpy image format
+        output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0
+        output_image = np.clip(output_image, 0, 255).astype(np.uint8)
+
+        # Free Memory
+        del image_tensor, output
+        torch.cuda.empty_cache()
+
+        return output_image
+
+# Create Gradio UI
+demo = gr.Interface(
+    fn=infer,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="numpy"),
+    title="Optimized Convolutional Reconstruction Model",
+    description="Upload an image to get the reconstructed output with reduced memory usage."
+)
+
+if __name__ == "__main__":
+    demo.launch()
+
+
+