daspartho / prompt-extend

extending stable diffusion prompts with suitable style cues using text generation

Prompt Extend

Text generation model for generating suitable style cues given the main idea for a prompt.

Play with it on HuggingFace Space.

Tokenizer and Model on HuggingFace Hub.

Usage

from transformers import pipeline

text_pipe = pipeline('text-generation', model='daspartho/prompt-extend')

prompt = "peaceful village landscape"
extended_prompt = text_pipe(prompt, num_return_sequences=1)[0]["generated_text"]
print(extended_prompt)

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("daspartho/prompt-extend")

model = AutoModelForCausalLM.from_pretrained("daspartho/prompt-extend")

Code

• dataset.ipynb: Creating a subset dataset of diffusiondb consisting of just unique prompts.
• tokenizer.ipynb: Training a new tokenizer on the dataset of stable diffusion prompts.
• model.ipynb: Training a GPT-2 model on the dataset of stable diffusion prompts using the trained tokenizer.
• inference.ipynb: Trying the prompts generated using the trained model with stable diffusion.
• app.ipynb: Gradio app for the model.

Examples

Contributing

If you want to contribute code, simply create a pull request. If you have an idea, create an issue and the developers will look into it!