From bc91ce11488e475827c9a2db5fcad1340c9ee924 Mon Sep 17 00:00:00 2001
From: Alex Tong <alexandertongdev@gmail.com>
Date: Wed, 1 Nov 2023 17:20:02 -0400
Subject: [PATCH] Updates to make torchcfm pip installable

---
 .gitignore       |  2 ++
 requirements.txt |  3 +--
 setup.py         | 25 +++++++++++++++++++++++--
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 58ea5eb..751869b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -171,3 +171,5 @@ slurm*.out
 *.jpg
 
 notebooks/figures/
+
+.DS_Store
diff --git a/requirements.txt b/requirements.txt
index 1816a9a..1b42df3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,9 +8,8 @@ scipy
 scikit-learn
 scprep
 scanpy
-timm
 torchdyn>=1.0.7   # 1.0.4 is broken on pypi
 pot
-torchdiffeq==0.2.3
+torchdiffeq
 absl-py
 clean-fid
diff --git a/setup.py b/setup.py
index 29ce294..ee660f3 100644
--- a/setup.py
+++ b/setup.py
@@ -4,15 +4,36 @@
 
 from setuptools import find_packages, setup
 
+install_requires = [
+    "torch>=1.11.0",
+    "torchvision>=0.11.0",
+    "lightning-bolts",
+    "matplotlib",
+    "numpy",
+    "scipy",
+    "scikit-learn",
+    "scprep",
+    "scanpy",
+    "torchdyn",
+    "pot",
+    "torchdiffeq",
+    "absl-py",
+    "clean-fid",
+]
+
 version_py = os.path.join(os.path.dirname(__file__), "torchcfm", "version.py")
 version = open(version_py).read().strip().split("=")[-1].replace('"', "").strip()
+readme = open("README.md").read()
 setup(
     name="torchcfm",
     version=version,
     description="Conditional Flow Matching for Fast Continuous Normalizing Flow Training.",
-    author="Alexander Tong",
+    author="Alexander Tong, Kilian Fatras",
     author_email="alexandertongdev@gmail.com",
     url="https://github.com/atong01/conditional-flow-matching",
-    install_requires=["torch", "pot", "numpy", "torchdyn"],
+    install_requires=install_requires,
+    license="MIT",
+    long_description=readme,
+    long_description_content_type="text/markdown",
     packages=find_packages(),
 )