Spaces:
Running
Running
| import gradio as gr | |
| #import peft | |
| import transformers | |
| import os | |
| import re | |
| device = "cpu" | |
| is_peft = False | |
| model_id = os.environ.get("MODEL_ID") or "treadon/prompt-fungineer-355M" | |
| auth_token = os.environ.get("HUB_TOKEN") or True | |
| print(f"Using model {model_id}.") | |
| if auth_token != True: | |
| print("Using auth token.") | |
| model = transformers.AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True,use_auth_token=auth_token) | |
| tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2") | |
| def format_prompt(prompt, enhancers=True, inspiration=False, negative_prompt=False): | |
| try: | |
| pattern = r"(BRF:|POS:|ENH:|INS:|NEG:) (.*?)(?= (BRF:|POS:|ENH:|INS:|NEG:)|$)" | |
| matches = re.findall(pattern, prompt) | |
| vals = {key: value.strip() for key, value,ex in matches} | |
| result = vals["POS:"] | |
| if enhancers: | |
| result += " " + vals["ENH:"] | |
| if inspiration: | |
| result += " " + vals["INS:"] | |
| if negative_prompt: | |
| result += "\n\n--no " + vals["NEG:"] | |
| return result | |
| except Exception as e: | |
| return "Failed to generate prompt." | |
| def generate_text(prompt, extra=False, top_k=100, top_p=0.95, temperature=0.85, enhancers = True, inpspiration = False , negative_prompt = False): | |
| if not prompt.startswith("BRF:"): | |
| prompt = "BRF: " + prompt | |
| if not extra: | |
| prompt = prompt + " POS:" | |
| model.eval() | |
| # SOFT SAMPLE | |
| inputs = tokenizer(prompt, return_tensors="pt").to(device) | |
| samples = [] | |
| try: | |
| for i in range(1): | |
| outputs = model.generate(**inputs, max_length=256, do_sample=True, top_k=top_k, top_p=top_p, temperature=temperature, num_return_sequences=4, pad_token_id=tokenizer.eos_token_id) | |
| for output in outputs: | |
| sample = tokenizer.decode(output, skip_special_tokens=True) | |
| sample = format_prompt(sample, enhancers, inpspiration, negative_prompt) | |
| samples.append(sample) | |
| except Exception as e: | |
| print(e) | |
| return samples | |
| with gr.Blocks() as fungineer: | |
| with gr.Row(): | |
| gr.Markdown("""# Midjourney / Dalle 2 / Stable Diffusion Prompt Generator | |
| This is the 355M parameter model. There is also a 7B parameter model that is much better but far slower (access coming soon). | |
| Just enter a basic prompt and the fungineering model will use its wildest imagination to expand the prompt in detail.""") | |
| with gr.Row(): | |
| with gr.Column(): | |
| base_prompt = gr.Textbox(lines=5, label="Base Prompt", placeholder="An astronaut in space", info="Enter a very simple prompt that will be fungineered into something exciting!") | |
| extra = gr.Checkbox(value=True, label="Extra Fungineer Imagination", info="If checked, the model will be allowed to go wild with its imagination.") | |
| with gr.Accordion("Advanced Generation Settings", open=False): | |
| top_k = gr.Slider( minimum=10, maximum=1000, value=100, label="Top K", info="Top K sampling") | |
| top_p = gr.Slider( minimum=0.1, maximum=1, value=0.95, step=0.01, label="Top P", info="Top P sampling") | |
| temperature = gr.Slider( minimum=0.1, maximum=1.2, value=0.85, step=0.01, label="Temperature", info="Temperature sampling. Higher values will make the model more creative") | |
| with gr.Accordion("Advanced Output Settings", open=False): | |
| enh = gr.Checkbox(value=True, label="Enhancers", info="Add image meta information such as lens type, shuffter speed, camera model, etc.") | |
| insp = gr.Checkbox(value=False, label="Inpsiration", info="Include inspirational photographers that are known for this type of photography. Sometimes random people will appear here, needs more training.") | |
| neg = gr.Checkbox(value=False, label="Negative Prompt", info="Include a negative prompt, more often used in Stable Diffusion. If you're a Stable Diffusion user, chances are you already have a better negative prompt you like to use.") | |
| with gr.Column(): | |
| outputs = [ | |
| gr.Textbox(lines=5, label="Fungineered Text 1"), | |
| gr.Textbox(lines=5, label="Fungineered Text 2"), | |
| gr.Textbox(lines=5, label="Fungineered Text 3"), | |
| gr.Textbox(lines=5, label="Fungineered Text 4"), | |
| ] | |
| inputs = [base_prompt, extra, top_k, top_p, temperature, enh, insp, neg] | |
| submit = gr.Button(label="Fungineer",variant="primary") | |
| submit.click(generate_text, inputs=inputs, outputs=outputs) | |
| fungineer.launch(enable_queue=True) | |