Update app.py

app.py

@@ -15,6 +15,11 @@ from transformers import AutoProcessor, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 from PIL import Image
 
+# os.makedirs("./hf_cache", exist_ok=True)
+# os.environ["HF_HOME"] = "./hf_cache"
+# os.environ["TRANSFORMERS_CACHE"] = "./hf_cache/transformers"
+# os.environ["HUGGINGFACE_HUB_CACHE"] = "./hf_cache/hub"
+
 # Add GPU decorator for Hugging Face Spaces
 try:
     from spaces import GPU
@@ -36,6 +41,11 @@ llava_model = None
 llava_processor = None
 stable_diffusion_pipeline = None
 
+# Set up the model directory
+MODEL_DIR = "./model"
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+# Update the model loading functions
 def load_llava_model():
     """Load LLaVA model for image captioning"""
     global llava_model, llava_processor
@@ -44,13 +54,16 @@ def load_llava_model():
         print("Loading LLaVA model for image analysis...")
         model_id = "llava-hf/llava-1.5-7b-hf"
         
-        # Load processor and model (with reduced precision for memory efficiency)
-        llava_processor = AutoProcessor.from_pretrained(model_id)
+        # Load processor and model with explicit cache directory
+        llava_processor = AutoProcessor.from_pretrained(
+            model_id,
+            cache_dir=os.path.join(MODEL_DIR, "llava_processor")
+        )
         llava_model = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
             device_map="auto",
-            cache_dir="./models"  # Add this line
+            cache_dir=os.path.join(MODEL_DIR, "llava_model")
         )
     
     return llava_model, llava_processor
@@ -67,8 +80,10 @@ def load_stable_diffusion_model():
         stable_diffusion_pipeline = StableDiffusionPipeline.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
-            safety_checker=None  # Disable safety checker for performance
+            safety_checker=None,  # Disable safety checker for performance
+            cache_dir=os.path.join(MODEL_DIR, "stable_diffusion")
         )
+
         
         # Move to GPU if available
         if torch.cuda.is_available():