Update app.py

app.py

@@ -1,53 +1,51 @@
 import gradio as gr
-import spaces  # ZeroGPU helper module
-from transformers import pipeline
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Preload the text-generation model on CPU at startup.
-# We load EleutherAI/gpt-j-6B on CPU (device=-1).
-generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=-1)
+# Load the MagicPrompt-Stable-Diffusion model and tokenizer
+model_name = "Gustavosta/MagicPrompt-Stable-Diffusion"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
 
-@spaces.GPU  # ZeroGPU will allocate the GPU only during this function call.
+@spaces.GPU
 def expand_prompt(prompt, num_variants=5, max_length=100):
     """
-    Given a basic prompt, generate `num_variants` expanded prompts using GPT-J-6B.
-    This function explicitly tokenizes the input with truncation (strategy 'longest_first'),
-    moves the input to GPU, generates output using the GPU, and then moves the model back to CPU.
+    Generate expanded prompts using a specialized model fine-tuned for Stable Diffusion.
     """
-    # Move model to GPU for generation.
-    generator.model.to("cuda")
+    # Move model to GPU
+    model.to("cuda")
     
-    # Explicitly tokenize the input with truncation.
-    inputs = generator.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length)
-    # Move inputs to GPU.
-    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    # Tokenize input prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     
-    # Generate text, explicitly setting pad_token_id to eos_token_id.
-    outputs = generator.model.generate(
+    # Generate multiple prompt variants
+    outputs = model.generate(
         **inputs,
         max_length=max_length,
-        num_return_sequences=num_variants,
         do_sample=True,
-        pad_token_id=generator.tokenizer.eos_token_id
+        num_return_sequences=num_variants,
+        pad_token_id=tokenizer.eos_token_id
     )
     
-    # Decode outputs.
-    expanded = [generator.tokenizer.decode(output, skip_special_tokens=True).strip() for output in outputs]
+    # Decode generated prompts
+    expanded_prompts = [tokenizer.decode(output, skip_special_tokens=True).strip() for output in outputs]
     
-    # Move model back to CPU.
-    generator.model.to("cpu")
+    # Move model back to CPU
+    model.to("cpu")
     
-    return "\n\n".join(expanded)
+    return "\n\n".join(expanded_prompts)
 
+# Create a Gradio Interface
 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 to receive 5 creative, expanded prompt variants. "
-        "The model (EleutherAI/gpt-j-6B) is preloaded on CPU at startup and moved to GPU (via ZeroGPU) for generation. "
-        "Input is tokenized with truncation enabled. Once generation is complete, the model is moved back to CPU. "
-        "Simply copy the output for use in your downstream image-generation pipeline."
+        "Enter a basic prompt and receive multiple expanded prompt variants optimized for Stable Diffusion. "
+        "This tool uses a specialized model fine-tuned on Stable Diffusion prompts. "
+        "Simply copy the output for use with your image-generation pipeline."
     )
 )