Update to latest ai-toolkit, torch, and cuda

replicate · Sep 7, 2024 · 18c4a1c · 18c4a1c
1 parent 3acf549
commit 18c4a1c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/ai-toolkit b/ai-toolkit
diff --git a/cog.yaml b/cog.yaml
@@ -3,12 +3,12 @@
 
 build:
   gpu: true
-  cuda: "12.1"
+  cuda: "12.4"
   python_version: "3.10.14"
   python_packages:
     # flux
-    - "torch==2.3.0"
-    - "torchvision==0.18.0"
+    - "torch==2.4.0"
+    - "torchvision==0.19.0"
     - "safetensors==0.4.4"
     - "diffusers==0.30.0"
     - "transformers==4.44.0"
@@ -56,6 +56,8 @@ build:
   run:
     - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64" && chmod +x /usr/local/bin/pget
     - pip install "git+https://github.com/Gothos/diffusers.git@flux-inpaint"
+    - "pip install --force-reinstall --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124"
+    #- "pip install --force-reinstall --pre torch==2.5.0.dev20240905+cu124 torchvision==0.20.0.dev20240905+cu124 --index-url https://download.pytorch.org/whl/nightly/cu124"
 
 # predict.py defines how predictions are run on your model
 predict: "predict.py:Predictor"
+2 −2		.github/ISSUE_TEMPLATE/bug_report.md
+1 −1		.github/ISSUE_TEMPLATE/config.yml
+125 −15		README.md
+ −		assets/lora_ease_ui.png
+96 −0		config/examples/modal/modal_train_lora_flux_24gb.yaml
+98 −0		config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml
+4 −0		config/examples/train_lora_flux_24gb.yaml
+4 −0		config/examples/train_lora_flux_schnell_24gb.yaml
+77 −27		extensions_built_in/sd_trainer/SDTrainer.py
+414 −0		flux_train_ui.py
+186 −39		jobs/process/BaseSDTrainProcess.py
+91 −75		notebooks/FLUX_1_dev_LoRA_Training.ipynb
+296 −0		notebooks/FLUX_1_schnell_LoRA_Training.ipynb
+3 −0		requirements.txt
+175 −0		run_modal.py
+32 −27		testing/test_bucket_dataloader.py
+16 −1		toolkit/config_modules.py
+19 −6		toolkit/custom_adapter.py
+7 −2		toolkit/dataloader_mixins.py
+245 −26		toolkit/ip_adapter.py
+3 −3		toolkit/lora_special.py
+10 −2		toolkit/models/DoRA.py
+6 −1		toolkit/models/ilora.py
+297 −39		toolkit/models/vd_adapter.py
+2 −1		toolkit/network_mixins.py
+13 −2		toolkit/samplers/custom_flowmatch_sampler.py
+1 −1		toolkit/scheduler.py
+6 −4		toolkit/stable_diffusion_model.py
+4 −6		toolkit/train_tools.py