import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer import spaces MODEL_NAME = "retowyss/PromptBridge-0.6b-Alpha" # Detect device DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32 DEFAULT_TEMPERATURE = 0.2 DEFAULT_MAX_TOKENS = 512 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=DTYPE, device_map="auto" if DEVICE == "cuda" else "cpu" ) model.eval() 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 = DEFAULT_TEMPERATURE, max_tokens: int = DEFAULT_MAX_TOKENS): """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").to(model.device) # Generate with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=max_tokens, do_sample=True, temperature=temperature, 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 expansion_examples = [ ["Expand", "woman, flowing red dress, standing, sunset beach", DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS], ["Expand", "man, leather jacket, motorcycle, urban alley", DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS], ["Expand", "dancer, ballet pose, stage lighting", DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS], ] compression_examples = [ ["Compress to Sentence", "A young woman with shoulder-length wavy auburn hair stands in a sunlit garden, her emerald green sundress flowing gently in the warm breeze. She wears delicate gold hoop earrings and a thin matching bracelet on her left wrist. Her bare feet rest on soft grass dotted with small white daisies. Behind her, climbing roses in shades of pink and coral wind up a weathered wooden trellis. The late afternoon sunlight filters through the leaves, casting dappled shadows across her face and creating a warm, golden glow that highlights her peaceful smile.", DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS], ["Compress to Keywords", "A young woman with shoulder-length wavy auburn hair stands in a sunlit garden, her emerald green sundress flowing gently in the warm breeze. She wears delicate gold hoop earrings and a thin matching bracelet on her left wrist. Her bare feet rest on soft grass dotted with small white daisies. Behind her, climbing roses in shades of pink and coral wind up a weathered wooden trellis. The late afternoon sunlight filters through the leaves, casting dappled shadows across her face and creating a warm, golden glow that highlights her peaceful smile.", DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS], ] # Create Gradio interface 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( choices=["Expand", "Compress to Sentence", "Compress to Keywords"], 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.0, maximum=1.5, value=DEFAULT_TEMPERATURE, step=0.1, label="Temperature" ) max_tokens = gr.Slider( minimum=64, maximum=1024, value=DEFAULT_MAX_TOKENS, step=32, 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, inputs=[mode, user_input, temperature, max_tokens], outputs=output, fn=generate_prompt, cache_examples=False ) with gr.Tab("Compression Examples"): gr.Examples( examples=compression_examples, inputs=[mode, user_input, temperature, max_tokens], outputs=output, 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, inputs=[mode, user_input, temperature, max_tokens], outputs=output ) if __name__ == "__main__": demo.launch()