LogicGoInfotechSpaces commited on
Commit
bdef219
·
1 Parent(s): 6bcce07

Update cache paths to use /tmp

Browse files
Files changed (3) hide show
  1. Dockerfile +2 -2
  2. app/colorize_model.py +5 -1
  3. app/main.py +11 -1
Dockerfile CHANGED
@@ -32,9 +32,9 @@ RUN mkdir -p /data/uploads /data/results
32
  # Note: Python code will create /data/hf_cache with proper permissions
33
  # Make Firebase credentials writing non-fatal (use || true to continue on error)
34
  RUN echo '#!/bin/sh' > /entrypoint.sh && \
35
- echo 'mkdir -p /data/uploads /data/results || true' >> /entrypoint.sh && \
36
  echo 'if [ -n "$FIREBASE_CREDENTIALS" ]; then' >> /entrypoint.sh && \
37
- echo ' printf "%s" "$FIREBASE_CREDENTIALS" > /data/firebase-adminsdk.json 2>/dev/null || true' >> /entrypoint.sh && \
38
  echo 'fi' >> /entrypoint.sh && \
39
  echo 'exec "$@"' >> /entrypoint.sh && \
40
  chmod +x /entrypoint.sh
 
32
  # Note: Python code will create /data/hf_cache with proper permissions
33
  # Make Firebase credentials writing non-fatal (use || true to continue on error)
34
  RUN echo '#!/bin/sh' > /entrypoint.sh && \
35
+ echo 'mkdir -p /tmp/colorize_uploads /tmp/colorize_results /tmp/hf_cache || true' >> /entrypoint.sh && \
36
  echo 'if [ -n "$FIREBASE_CREDENTIALS" ]; then' >> /entrypoint.sh && \
37
+ echo ' printf "%s" "$FIREBASE_CREDENTIALS" > /tmp/firebase-adminsdk.json 2>/dev/null || true' >> /entrypoint.sh && \
38
  echo 'fi' >> /entrypoint.sh && \
39
  echo 'exec "$@"' >> /entrypoint.sh && \
40
  chmod +x /entrypoint.sh
app/colorize_model.py CHANGED
@@ -96,10 +96,14 @@ class ColorizeModel:
96
 
97
  def _load_caption_model(self) -> None:
98
  logger.info("Loading BLIP captioning model for prompt generation...")
99
- self.caption_processor = BlipProcessor.from_pretrained(self.CAPTION_MODEL)
 
 
 
100
  self.caption_model = BlipForConditionalGeneration.from_pretrained(
101
  self.CAPTION_MODEL,
102
  torch_dtype=self.dtype if self.device.type == "cuda" else torch.float32,
 
103
  ).to(self.device)
104
 
105
  def caption_image(self, image: Image.Image) -> str:
 
96
 
97
  def _load_caption_model(self) -> None:
98
  logger.info("Loading BLIP captioning model for prompt generation...")
99
+ self.caption_processor = BlipProcessor.from_pretrained(
100
+ self.CAPTION_MODEL,
101
+ cache_dir=self.cache_dir
102
+ )
103
  self.caption_model = BlipForConditionalGeneration.from_pretrained(
104
  self.CAPTION_MODEL,
105
  torch_dtype=self.dtype if self.device.type == "cuda" else torch.float32,
106
+ cache_dir=self.cache_dir
107
  ).to(self.device)
108
 
109
  def caption_image(self, image: Image.Image) -> str:
app/main.py CHANGED
@@ -21,6 +21,16 @@ import io
21
  from app.colorize_model import ColorizeModel
22
  from app.config import settings
23
 
 
 
 
 
 
 
 
 
 
 
24
  # Configure logging
25
  logging.basicConfig(
26
  level=logging.INFO,
@@ -45,7 +55,7 @@ app.add_middleware(
45
  )
46
 
47
  # Initialize Firebase Admin SDK
48
- firebase_cred_path = os.getenv("FIREBASE_CREDENTIALS_PATH", "colorize-662df-firebase-adminsdk-fbsvc-e080668793.json")
49
  if os.path.exists(firebase_cred_path):
50
  try:
51
  cred = credentials.Certificate(firebase_cred_path)
 
21
  from app.colorize_model import ColorizeModel
22
  from app.config import settings
23
 
24
+ # Ensure HF caches and temporary directories use /tmp (writable in Spaces)
25
+ os.environ.setdefault("HF_HOME", "/tmp/hf_cache")
26
+ os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf_cache")
27
+ os.environ.setdefault("HF_HUB_CACHE", "/tmp/hf_cache")
28
+ os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/hf_cache")
29
+ os.environ.setdefault("XDG_CACHE_HOME", "/tmp/hf_cache")
30
+ Path("/tmp/hf_cache").mkdir(parents=True, exist_ok=True)
31
+ Path("/tmp/colorize_uploads").mkdir(parents=True, exist_ok=True)
32
+ Path("/tmp/colorize_results").mkdir(parents=True, exist_ok=True)
33
+
34
  # Configure logging
35
  logging.basicConfig(
36
  level=logging.INFO,
 
55
  )
56
 
57
  # Initialize Firebase Admin SDK
58
+ firebase_cred_path = os.getenv("FIREBASE_CREDENTIALS_PATH", "/tmp/firebase-adminsdk.json")
59
  if os.path.exists(firebase_cred_path):
60
  try:
61
  cred = credentials.Certificate(firebase_cred_path)