From c8fb30aaeeadb49041d551d36d3193f5ae1546ff Mon Sep 17 00:00:00 2001
From: zcr <cgzhou@aidlab.hk>
Date: Wed, 18 Mar 2026 11:18:49 +0800
Subject: [PATCH] 1

---
 .gitignore            | 3 ++-
 app/litserve_serve.py | 1 +
 docker-compose.yml    | 1 +
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 3dc0ecf..b7337c7 100755
--- a/.gitignore
+++ b/.gitignore
@@ -230,4 +230,5 @@ $RECYCLE.BIN/
 # End of https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
 
 output/
-*.png
\ No newline at end of file
+*.png
+checkpoint/
\ No newline at end of file
diff --git a/app/litserve_serve.py b/app/litserve_serve.py
index c57d2e6..fa7e6cf 100755
--- a/app/litserve_serve.py
+++ b/app/litserve_serve.py
@@ -37,6 +37,7 @@ class FluxKleinAPI(ls.LitAPI):
 
         self.pipe = Flux2KleinPipeline.from_pretrained(
             self.repo_id,
+            cache_dir="./checkpoint/",
             torch_dtype=self.dtype
         )
         self.pipe.to(device)
diff --git a/docker-compose.yml b/docker-compose.yml
index ba8ee59..e7cb752 100755
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -7,6 +7,7 @@ services:
     working_dir: /app
     volumes:
       - ./app:/app/app
+      - ./checkpoint:/app/checkpoint
       - /etc/localtime:/etc/localtime:ro
     ports:
       - "14202:8451"