Update app.py

app.py

@@ -1,16 +1,29 @@
 import gradio as gr
 from transformers import pipeline
+import torch
 
-# Load the larger text-generation model that uses GPU.
-# Here we use EleutherAI/gpt-j-6B: https://huggingface.co/EleutherAI/gpt-j-6B
-# Setting device=0 tells the pipeline to use GPU 0.
-generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=0)
+# Global generator variable; we'll load it lazily.
+generator = None
+
+def get_generator():
+    global generator
+    if generator is None:
+        try:
+            # 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
 
 def expand_prompt(prompt, num_variants=5, max_length=100):
-    """
-    Given a basic prompt, generate `num_variants` expanded prompts using GPT-J-6B.
-    """
-    outputs = generator(prompt, max_length=max_length, num_return_sequences=num_variants, do_sample=True)
+    # Lazy load the model when a prompt is submitted.
+    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]
     return "\n\n".join(expanded)
 
@@ -21,8 +34,8 @@ 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 on an A100 GPU for fast, expressive prompt expansion. "
-        "Simply copy the output for use with your downstream image-generation pipeline."
+        "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."
     )
 )