Update app.py

app.py

@@ -3,6 +3,9 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import spaces
 
+import gc 
+import torch
+
 # Create the necessary directories
 os.makedirs('.gradio/cached_examples/17', exist_ok=True)
 
@@ -34,6 +37,9 @@ def load_model(model_name):
         top_k=50,        # Control diversity
         top_p=0.95       # Control diversity
     )
+    del model 
+    del tokenizer
+    
     return generator
 
 @spaces.GPU
@@ -51,18 +57,15 @@ def generate_kids_story(character, setting, language):
 
     messages = [{"role": "user", "content": prompt}]
     output = generator(messages)
-    return output[0]["generated_text"]
 
-import gc 
-import torch
-# Delete model and associated objects
-del model 
-del tokenizer 
-del generator
-# Run garbage collection
-gc.collect ()
-# Empty CUDA cache
-torch.cuda.empty_cache()
+    # Delete model and associated objects 
+    del generator
+    # Run garbage collection
+    gc.collect ()
+    # Empty CUDA cache
+    torch.cuda.empty_cache()
+    
+    return output[0]["generated_text"]
 
 # Create Gradio interface
 demo = gr.Interface(