Make it work

README.md

@@ -1,12 +1,9 @@
 ---
-title: Prompt Fungineer 355M
-emoji: 🐠
+title: Midjourney / Dalle 2 / Stable Diffusion Prompt Generator (Prompt Fungineer 355M)
+emoji: 🧙🏼‍♂️
 colorFrom: red
 colorTo: yellow
 sdk: gradio
 sdk_version: 3.24.1
 app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+pinned: false

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 #import peft
 import transformers
 import os
+import re
 
 device = "cpu"
 is_peft = False
@@ -18,25 +19,92 @@ auth_token = os.environ.get("hub_token") or True
 model = transformers.AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True,use_auth_token=auth_token)
 tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
 
-def generate_text(prompt):
+
+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=100, top_p=0.95, temperature=0.85, num_return_sequences=4, pad_token_id=tokenizer.eos_token_id)
+            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
 
-iface = gr.Interface(fn=generate_text, inputs="text", outputs=("text","text","text","text") )
-iface.launch()
+
+# inputs = [
+#     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!"),
+#     gr.Checkbox(value=True, label="Extra Fungineer Imagination", info="If checked, the model will be allowed to go wild with its imagination."),
+#     gr.Slider( minimum=10, maximum=1000, value=100, label="Top K", info="Top K sampling"),
+#     gr.Slider( minimum=0.1, maximum=1, value=0.95, step=0.01, label="Top P", info="Top P sampling"),
+#     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"),
+# ]
+
+# iface = gr.Interface(fn=generate_text, inputs=inputs, outputs=["text","text","text","text"] )
+
+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):
+                gr.Checkbox(value=True, label="Enhancers", info="Add image meta information such as lens type, shuffter speed, camera model, etc.")
+                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.")
+                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]
+
+
+    submit = gr.Button(label="Fungineer",variant="primary")
+    submit.click(generate_text, inputs=inputs, outputs=outputs)
+
+fungineer.launch(enable_queue=True)
+