flake/overlays/python/default.nix

lib: {
  fixPackages = final: prev: let
    relaxProtobuf = pkg:
      pkg.overrideAttrs (old: {
        nativeBuildInputs = old.nativeBuildInputs ++ [final.pythonRelaxDepsHook];
        pythonRelaxDeps = ["protobuf"];
      });
  in {
    invisible-watermark = prev.invisible-watermark.overridePythonAttrs {
      pythonImportsCheck = [];
    };
    torchsde = prev.torchsde.overridePythonAttrs {
      doCheck = false;
      pythonImportsCheck = [];
    };
    pytorch-lightning = relaxProtobuf prev.pytorch-lightning;
    wandb = relaxProtobuf (prev.wandb.overridePythonAttrs {
      doCheck = false;
    });
    anyio = prev.anyio.overridePythonAttrs {
      doCheck = false;
      dontUsePytestCheck = true;
    };
    mocket = prev.mocket.overridePythonAttrs {
      doCheck = false;
      dontUsePytestCheck = true;
    };
    markdown-it-py = prev.markdown-it-py.overrideAttrs (old: {
      nativeBuildInputs = old.nativeBuildInputs ++ [final.pythonRelaxDepsHook];
      pythonRelaxDeps = ["linkify-it-py"];
      passthru =
        old.passthru
        // {
          optional-dependencies = with final; {
            linkify = [linkify-it-py];
            plugins = [mdit-py-plugins];
          };
        };
    });
    filterpy = prev.filterpy.overrideAttrs (old: {
      doInstallCheck = false;
    });
    shap = prev.shap.overrideAttrs (old: {
      doInstallCheck = false;
      propagatedBuildInputs = old.propagatedBuildInputs ++ [final.packaging];
      pythonImportsCheck = ["shap"];

      meta =
        old.meta
        // {
          broken = false;
        };
    });
    streamlit = let
      streamlit = final.callPackage (final.pkgs.path + "/pkgs/applications/science/machine-learning/streamlit") {
        protobuf3 = final.protobuf;
      };
    in
      final.toPythonModule (relaxProtobuf streamlit);
    opencv-python-headless = final.opencv-python;
    opencv-python = final.opencv4;
  };

  torchRocm = final: prev: {
    torch = prev.torch.override {
      magma = prev.pkgs.magma-hip;
      cudaSupport = false;
      rocmSupport = true;
    };
    torchvision = prev.torchvision.overridePythonAttrs (old: {
      patches = (old.patches or []) ++ [./torchvision/fix-rocm-build.patch];
    });
  };

  torchCuda = final: prev: {
    torch = prev.torch.override {
      magma = prev.pkgs.magma-cuda-static;
      cudaSupport = true;
      rocmSupport = false;
    };
  };

  bitsAndBytesOldGpu = final: prev: {
    bitsandbytes = prev.bitsandbytes.overridePythonAttrs (old: {
      preBuild = old.preBuild + " cuda${final.torch.cudaPackages.cudaMajorVersion}x_nomatmul";
    });
  };
}