Update app.py

app.py

@@ -1,27 +1,31 @@
 import gradio as gr
+import spaces  # Import ZeroGPU's helper module
 from transformers import pipeline
 import torch
 
-# Global generator variable; we'll load it lazily.
+# Global generator variable; load lazily.
 generator = None
 
 def get_generator():
     global generator
     if generator is None:
         try:
-            # If GPU is available, load on GPU (device=0).
+            # If GPU is available, load on GPU (device=0)
             if torch.cuda.is_available():
                 generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=0)
             else:
                 generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=-1)
         except Exception as e:
-            # If any error occurs, fallback to CPU
             print("Error loading model on GPU, falling back to CPU:", e)
             generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=-1)
     return generator
 
+@spaces.GPU  # This decorator ensures ZeroGPU allocates a GPU when the function is called.
 def expand_prompt(prompt, num_variants=5, max_length=100):
-    # Lazy load the model when a prompt is submitted.
+    """
+    Given a basic prompt, generate `num_variants` expanded prompts using GPT-J-6B.
+    The GPU is only engaged during this function call.
+    """
     gen = get_generator()
     outputs = gen(prompt, max_length=max_length, num_return_sequences=num_variants, do_sample=True)
     expanded = [out["generated_text"].strip() for out in outputs]
@@ -34,8 +38,9 @@ iface = gr.Interface(
     title="Prompt Expansion Generator",
     description=(
         "Enter a basic prompt and receive 5 creative, expanded prompt variants. "
-        "This tool leverages the EleutherAI/gpt-j-6B model and defers loading it until the first prompt is received—"
-        "letting ZeroGPU initialize properly. Simply copy the output for use with your downstream image-generation pipeline."
+        "This tool leverages the EleutherAI/gpt-j-6B model on an A100 GPU via ZeroGPU. "
+        "The GPU is only allocated when a prompt is submitted, ensuring proper ZeroGPU initialization. "
+        "Simply copy the output for use with your downstream image-generation pipeline."
     )
 )