Update app.py

app.py

@@ -1,36 +1,27 @@
 import gradio as gr
-import spaces  # Import ZeroGPU's helper module
+import spaces  # ZeroGPU helper module
 from transformers import pipeline
-import torch
 
-# Global generator variable; load lazily.
-generator = None
+# Preload the text-generation model on CPU at startup.
+# Model: EleutherAI/gpt-j-6B (https://huggingface.co/EleutherAI/gpt-j-6B)
+# We load on CPU (device=-1) so that initialization is done before the GUI is up.
+generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=-1)
 
-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:
-            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.
+@spaces.GPU  # This decorator ensures ZeroGPU allocates a GPU only during this function call.
 def expand_prompt(prompt, num_variants=5, max_length=100):
     """
-    Given a basic prompt, generate `num_variants` expanded prompts using GPT-J-6B.
-    The GPU is only engaged during this function call.
+    Given a basic prompt, generate `num_variants` expanded prompt variants.
+    Before generation, the model is moved to GPU (A100), and after generation it's moved back to CPU.
     """
-    gen = get_generator()
-    outputs = gen(prompt, max_length=max_length, num_return_sequences=num_variants, do_sample=True)
+    # Move the model to GPU for generation.
+    generator.model.to("cuda")
+    outputs = generator(prompt, max_length=max_length, num_return_sequences=num_variants, do_sample=True)
+    # Move the model back to CPU after generation.
+    generator.model.to("cpu")
     expanded = [out["generated_text"].strip() for out in outputs]
     return "\n\n".join(expanded)
 
+# Create a Gradio Interface
 iface = gr.Interface(
     fn=expand_prompt,
     inputs=gr.Textbox(lines=2, placeholder="Enter your basic prompt here...", label="Basic Prompt"),
@@ -38,9 +29,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 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."
+        "The model (EleutherAI/gpt-j-6B) is preloaded on CPU at startup and then moved to GPU (via ZeroGPU) only "
+        "when a prompt is submitted. Simply copy the output for use with your downstream image-generation pipeline."
     )
 )