Remove editor module

src/editor/models/__init__.py

@@ -1,76 +0,0 @@
-from typing import Any, Dict, Tuple
-from .v1 import HistogramRestorationNet as v1
-from .simple_cnn import SimpleCNN
-from .residual import Residual
-from .residual3 import Residual3
-from .dummy import Dummy
-import torch
-import torch.nn as nn
-from pathlib import Path
-import logging
-import json
-
-
-MODELS = {
-    # "v1": v1,
-    "Dummy": Dummy,
-    "SimpleCNN": SimpleCNN,
-    "Residual": Residual,
-    "Residual3": Residual3,
-}
-
-
-def create_model(type: str, bin_count: int, device: torch.device) -> nn.Module:
-    return MODELS[type](bin_count).to(device)
-
-
-def save_model(model: nn.Module, hyperparameters: Dict[str, Any], path: Path):
-    model_path = path.with_suffix(".pth")
-    params_path = path.with_suffix(".json")
-
-    logging.info(f"Saving model to {model_path}")
-    logging.info(f"Parameter count: {sum(p.numel() for p in model.parameters())}")
-    with open(model_path, "wb") as f:
-        torch.save(model.state_dict(), f)
-    with open(params_path, "w") as f:
-        json.dump(hyperparameters, f, indent=2)
-
-
-def load_model(path: Path, device: torch.device) -> Tuple[nn.Module, Dict[str, Any]]:
-    logging.info(f"Loading model from {path}")
-
-    params_path = path.with_suffix(".json")
-    with open(params_path) as f:
-        hyperparameters = json.load(f)
-    logging.info(f"Hyperparameters: {hyperparameters}")
-
-    model_path = path.with_suffix(".pth")
-    model = create_model(
-        type=hyperparameters["model_type"],
-        bin_count=hyperparameters["bin_count"],
-    ).to(device)
-    model.load_state_dict(torch.load(model_path))
-    model.eval()
-    logging.info(f"Parameter count: {sum(p.numel() for p in model.parameters())}")
-
-    return model, hyperparameters
-
-
-def test_models():
-    for model_name, model_constructor in MODELS.items():
-        logging.info(f"Testing model {model_name}")
-        _test_network_dimensions(model_constructor)
-
-
-def _test_network_dimensions(constructor):
-    for bin_count in [16, 32, 64]:
-        model = constructor()
-
-        # Create a dummy input tensor of the correct shape, the mini-batch size is 4
-        input_tensor = torch.rand(4, 1, bin_count, bin_count, bin_count)
-
-        output = model(input_tensor)
-        assert (
-            input_tensor.shape == output.shape
-        ), f"Expected output shape {input_tensor.shape}, but got {output.shape}"
-    logging.info("Test passed! Output shape matches input shape.")