Spaces:
Sleeping
Sleeping
File size: 1,782 Bytes
368407f e94cd94 368407f e94cd94 368407f e94cd94 368407f e94cd94 368407f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import gradio as gr
from transformers import pipeline
import torch
# 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):
# 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)
iface = gr.Interface(
fn=expand_prompt,
inputs=gr.Textbox(lines=2, placeholder="Enter your basic prompt here...", label="Basic Prompt"),
outputs=gr.Textbox(lines=10, label="Expanded Prompts"),
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."
)
)
if __name__ == "__main__":
iface.launch()
|