Update app.py

app.py

@@ -1,27 +1,76 @@
-import gradio as gr
+# import gradio as gr
+# import torch
+# from PIL import Image
+# from model import CRM
+# from inference import generate3d
+# import numpy as np
+
+# # Load model
+# crm_path = "CRM.pth"  # Make sure the model is uploaded to the Space
+# model = CRM(torch.load(crm_path, map_location="cpu"))
+# model = model.to("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# def generate_3d(image_path, seed=1234, scale=5.5, step=30):
+#     image = Image.open(image_path).convert("RGB")
+#     np_img = np.array(image)
+#     glb_path = generate3d(model, np_img, np_img, "cuda:0" if torch.cuda.is_available() else "cpu")
+#     return glb_path
+
+# iface = gr.Interface(
+#     fn=generate_3d,
+#     inputs=gr.Image(type="filepath"),
+#     outputs=gr.Model3D(),
+#     title="Convolutional Reconstruction Model (CRM)",
+#     description="Upload an image to generate a 3D model."
+# )
+
+# iface.launch()
+import os
 import torch
-from PIL import Image
-from model import CRM
-from inference import generate3d
-import numpy as np
-
-# Load model
-crm_path = "CRM.pth"  # Make sure the model is uploaded to the Space
-model = CRM(torch.load(crm_path, map_location="cpu"))
-model = model.to("cuda:0" if torch.cuda.is_available() else "cpu")
-
-def generate_3d(image_path, seed=1234, scale=5.5, step=30):
-    image = Image.open(image_path).convert("RGB")
-    np_img = np.array(image)
-    glb_path = generate3d(model, np_img, np_img, "cuda:0" if torch.cuda.is_available() else "cpu")
-    return glb_path
-
-iface = gr.Interface(
-    fn=generate_3d,
-    inputs=gr.Image(type="filepath"),
-    outputs=gr.Model3D(),
+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 3D model."
+    description="Upload an image to generate a reconstructed output."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()