fix loading issue

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gc
 
@@ -24,9 +24,25 @@ class SpacesHaikuGenerator:
             "Reflective": "Write a contemplative haiku about {name}, focusing on {traits}"
         }
 
-    @st.cache_resource
     def load_model(self, model_name):
-        """Load model with caching for Streamlit."""
+        """Load model with proper caching."""
+
+        @st.cache_resource
+        def _load_model_cached(_model_name):
+            tokenizer = AutoTokenizer.from_pretrained(
+                self.models[_model_name],
+                trust_remote_code=True
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                self.models[_model_name],
+                trust_remote_code=True,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True
+            )
+            if torch.cuda.is_available():
+                model = model.to("cuda")
+            return model, tokenizer
+
         if self.current_model != model_name:
             # Clear previous model
             if self.loaded_model is not None:
@@ -35,22 +51,10 @@ class SpacesHaikuGenerator:
                 torch.cuda.empty_cache()
                 gc.collect()
 
-            # Load new model
-            self.loaded_tokenizer = AutoTokenizer.from_pretrained(
-                self.models[model_name],
-                trust_remote_code=True
-            )
-            self.loaded_model = AutoModelForCausalLM.from_pretrained(
-                self.models[model_name],
-                trust_remote_code=True,
-                torch_dtype=torch.float16,
-                low_cpu_mem_usage=True
-            )
+            # Load new model using cached function
+            self.loaded_model, self.loaded_tokenizer = _load_model_cached(model_name)
             self.current_model = model_name
 
-            if torch.cuda.is_available():
-                self.loaded_model = self.loaded_model.to("cuda")
-
     def generate_haiku(self, name, traits, model_name, style):
         """Generate a free-form haiku using the selected model."""
         self.load_model(model_name)
@@ -103,7 +107,11 @@ def main():
     st.write("Create unique AI-generated haikus about characters")
 
     # Initialize generator
-    generator = SpacesHaikuGenerator()
+    @st.cache_resource
+    def get_generator():
+        return SpacesHaikuGenerator()
+
+    generator = get_generator()
 
     # Input fields
     col1, col2 = st.columns([1, 2])