Upload 2 files

app.py

@@ -446,7 +446,6 @@ class ModelState:
     """Global state for loaded model."""
 
     model: Optional[FeatureFusionModel] = None
-    device: str = "cpu"
 
 
 state = ModelState()
@@ -457,15 +456,13 @@ def load_model() -> str:
     if state.model is not None:
         return "Model already loaded"
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
     try:
         checkpoint_path = hf_hub_download(
             repo_id="marduk-ra/MFIR",
             filename="temporal_fusion_model.pth",
         )
 
-        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
 
         config_dict = ckpt["config"]
         if isinstance(config_dict, dict):
@@ -482,13 +479,11 @@ def load_model() -> str:
         config = FeatureFusionConfig.from_dict(config_dict)
         model = FeatureFusionModel(config)
         model.load_state_dict_with_compatibility(ckpt["state_dict"])
-        model.to(device)
         model.eval()
 
         state.model = model
-        state.device = device
 
-        return f"Model loaded on {device}"
+        return "Model loaded"
 
     except Exception as e:
         return f"Error loading model: {e}"
@@ -609,8 +604,6 @@ def process_images(
         if frames is None:
             return None, "Failed to preprocess images"
 
-        frames = frames.to(state.device)
-
         ref_idx = min(ref_frame, frames.shape[1] - 1)
 
         with torch.no_grad():
@@ -619,10 +612,6 @@ def process_images(
         output = result["output"]
         output_pil = postprocess_output(output, preprocess_info)
 
-        del result, output, frames
-        if state.device == "cuda":
-            torch.cuda.empty_cache()
-
         return output_pil, f"Processed {len(images)} frames (reference: frame {ref_idx})"
 
     except Exception as e:

requirements.txt

@@ -1,5 +1,5 @@
-torch>=2.1.0
-torchvision>=0.16.0
+torch
+torchvision
 huggingface_hub
 numpy
 Pillow