Update app.py

app.py

@@ -1,47 +1,57 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
 
 MODEL_NAME = "retowyss/PromptBridge-0.6b-Alpha"
 
-# Load model and tokenizer
+# Detect device
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
+
+print(f"Using device: {DEVICE}")
 print("Loading model...")
+
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+
+# Load model with appropriate settings
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     trust_remote_code=True,
-    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-    device_map="cpu"
+    torch_dtype=DTYPE,
+    device_map="auto" if DEVICE == "cuda" else "cpu"
 )
 model.eval()
-print("Model loaded!")
+print(f"Model loaded on {DEVICE}!")
 
+@spaces.GPU(duration=60)  # Allocate GPU for 60 seconds
 def generate_prompt(mode: str, user_prompt: str, temperature: float = 0.7, max_tokens: int = 512):
     """Generate prompt transformation."""
-    
+
     # Map mode to system prompt
     system_prompts = {
         "Expand": "Expand the prompt.",
         "Compress to Sentence": "Compress the prompt into one sentence.",
         "Compress to Keywords": "Compress the prompt into keyword format."
     }
-    
+
     system_prompt = system_prompts[mode]
-    
+
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt}
     ]
-    
+
     # Apply chat template
     text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    
-    inputs = tokenizer(text, return_tensors="pt")
-    
+
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
     # Generate
     with torch.no_grad():
         outputs = model.generate(
@@ -53,13 +63,13 @@ def generate_prompt(mode: str, user_prompt: str, temperature: float = 0.7, max_t
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
         )
-    
+
     # Decode only the new tokens
     response = tokenizer.decode(
         outputs[0][inputs['input_ids'].shape[1]:],
         skip_special_tokens=True
     )
-    
+
     return response.strip()
 
 # Example prompts
@@ -78,19 +88,19 @@ compression_examples = [
 with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
     gr.Markdown("""
     # 🌉 PromptBridge-0.6b-Alpha
-    
+
     A specialized model for bidirectional prompt transformation for text-to-image generation.
-    
+
     **Trained exclusively on prompts featuring single, adult humanoid subjects.**
-    
+
     ### Modes:
     - **Expand**: Convert brief keywords into detailed image generation prompts
     - **Compress to Sentence**: Condense detailed prompts into a single flowing sentence
     - **Compress to Keywords**: Convert prompts into comma-separated keywords
-    
+
     ⚠️ **Note**: May generate sensitive content (PG to R-rated). You are responsible for the output.
     """)
-    
+
     with gr.Row():
         with gr.Column():
             mode = gr.Radio(
@@ -98,13 +108,13 @@ with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
                 value="Expand",
                 label="Mode"
             )
-            
+
             user_input = gr.Textbox(
                 lines=5,
                 placeholder="Enter your prompt here...",
                 label="Input Prompt"
             )
-            
+
             with gr.Row():
                 temperature = gr.Slider(
                     minimum=0.1,
@@ -113,7 +123,7 @@ with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
                     step=0.1,
                     label="Temperature"
                 )
-                
+
                 max_tokens = gr.Slider(
                     minimum=64,
                     maximum=1024,
@@ -121,18 +131,18 @@ with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
                     step=64,
                     label="Max Tokens"
                 )
-            
+
             submit_btn = gr.Button("Generate", variant="primary")
-        
+
         with gr.Column():
             output = gr.Textbox(
                 lines=15,
                 label="Output"
             )
-    
+
     # Examples
     gr.Markdown("### Examples")
-    
+
     with gr.Tab("Expansion Examples"):
         gr.Examples(
             examples=expansion_examples,
@@ -141,7 +151,7 @@ with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
             fn=generate_prompt,
             cache_examples=False
         )
-    
+
     with gr.Tab("Compression Examples"):
         gr.Examples(
             examples=compression_examples,
@@ -150,22 +160,22 @@ with gr.Blocks(title="PromptBridge-0.6b-Alpha") as demo:
             fn=generate_prompt,
             cache_examples=False
         )
-    
+
     gr.Markdown("""
     ### About
-    
+
     **Model**: [retowyss/PromptBridge-0.6b-Alpha](https://huggingface.co/retowyss/PromptBridge-0.6b-Alpha)
-    
+
     **Training Data**: ~300K synthetic prompt pairs (PG to R-rated, X-rated content removed)
-    
+
     **Limitations**:
     - Optimized for human subjects only (performance on other subjects untested)
     - May generate prompts exceeding typical image model token limits
     - Cannot perform general instruction-following or reasoning tasks
-    
+
     **License**: Apache 2.0
     """)
-    
+
     # Connect the button
     submit_btn.click(
         fn=generate_prompt,