app.py

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()