Update app.py

app.py

@@ -25,51 +25,112 @@
 # )
 
 # iface.launch()
-import os
+
+
+#############2nd################3
+# import os
+# import torch
+# import gradio as gr
+# from huggingface_hub import hf_hub_download
+# from model import CRM  # Make sure this matches your model file structure
+
+# # Define model details
+# REPO_ID = "Mariam-Elz/CRM"  # Hugging Face model repo
+# MODEL_FILES = {
+#     "ccm-diffusion": "ccm-diffusion.pth",
+#     "pixel-diffusion": "pixel-diffusion.pth",
+#     "CRM": "CRM.pth"
+# }
+# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# # Download models from Hugging Face if not already present
+# MODEL_DIR = "./models"
+# os.makedirs(MODEL_DIR, exist_ok=True)
+
+# for name, filename in MODEL_FILES.items():
+#     model_path = os.path.join(MODEL_DIR, filename)
+#     if not os.path.exists(model_path):
+#         print(f"Downloading {filename}...")
+#         hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=MODEL_DIR)
+
+# # Load the model
+# print("Loading CRM Model...")
+# model = CRM()
+# model.load_state_dict(torch.load(os.path.join(MODEL_DIR, MODEL_FILES["CRM"]), map_location=DEVICE))
+# model.to(DEVICE)
+# model.eval()
+# print("✅ Model Loaded Successfully!")
+
+# # Define Gradio Interface
+# def predict(input_image):
+#     with torch.no_grad():
+#         output = model(input_image.to(DEVICE))  # Modify based on model input format
+#     return output.cpu()
+
+# demo = gr.Interface(
+#     fn=predict,
+#     inputs=gr.Image(type="pil"),
+#     outputs=gr.Image(type="pil"),
+#     title="Convolutional Reconstruction Model (CRM)",
+#     description="Upload an image to generate a reconstructed output."
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+########################3rd######################3
+
 import torch
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from model import CRM  # Make sure this matches your model file structure
-
-# Define model details
-REPO_ID = "Mariam-Elz/CRM"  # Hugging Face model repo
-MODEL_FILES = {
-    "ccm-diffusion": "ccm-diffusion.pth",
-    "pixel-diffusion": "pixel-diffusion.pth",
-    "CRM": "CRM.pth"
+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",
 }
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Download models from Hugging Face if not already present
-MODEL_DIR = "./models"
-os.makedirs(MODEL_DIR, exist_ok=True)
+os.makedirs("models", exist_ok=True)
 
-for name, filename in MODEL_FILES.items():
-    model_path = os.path.join(MODEL_DIR, filename)
-    if not os.path.exists(model_path):
+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}...")
-        hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=MODEL_DIR)
-
-# Load the model
-print("Loading CRM Model...")
-model = CRM()
-model.load_state_dict(torch.load(os.path.join(MODEL_DIR, MODEL_FILES["CRM"]), map_location=DEVICE))
-model.to(DEVICE)
-model.eval()
-print("✅ Model Loaded Successfully!")
-
-# Define Gradio Interface
-def predict(input_image):
+        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():
-        output = model(input_image.to(DEVICE))  # Modify based on model input format
-    return output.cpu()
+        # 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=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Convolutional Reconstruction Model (CRM)",
-    description="Upload an image to generate a reconstructed output."
+    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__":