Update src/app_hf_space_optimized.py

src/app_hf_space_optimized.py

@@ -35,11 +35,12 @@ def clear_torch():
 # --- Step 1: Generate JSON Story ---
 def generate_story(prompt: str, num_scenes: int):
     st.info("🧠 Generating story...")
-    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
+    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID, cache_dir=CACHE_DIR)
     model = AutoModelForCausalLM.from_pretrained(
         LLM_MODEL_ID,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
+        device_map="auto",
+        cache_dir=CACHE_DIR
     )
 
     sys_prompt = (
@@ -69,7 +70,8 @@ def generate_images(scenes):
     st.info("🎨 Generating images...")
     pipe = StableDiffusionPipeline.from_pretrained(
         IMG_MODEL_ID,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        cache_dir=CACHE_DIR
     )
     pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     images = []
@@ -83,9 +85,9 @@ def generate_images(scenes):
 # --- Step 3: Generate TTS ---
 def generate_audios(scenes):
     st.info("🔊 Generating audio...")
-    tts = ParlerTTSForConditionalGeneration.from_pretrained(TTS_MODEL_ID, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_ID)
-    desc_tokenizer = AutoTokenizer.from_pretrained(tts.config.text_encoder._name_or_path)
+    tts = ParlerTTSForConditionalGeneration.from_pretrained(TTS_MODEL_ID, device_map="auto", cache_dir=CACHE_DIR)
+    tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_ID, cache_dir=CACHE_DIR)
+    desc_tokenizer = AutoTokenizer.from_pretrained(tts.config.text_encoder._name_or_path, cache_dir=CACHE_DIR)
 
     audio_paths = []
     for i, scene in enumerate(scenes):