Update app.py

app.py

@@ -5,6 +5,25 @@ import random
 import os
 from PIL import Image
 
+from datasets import load_dataset
+
+dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts")
+prompt_df = dataset["train"].to_pandas()
+
+#DEFAULT_MODEL = "stabilityai/stable-diffusion-2-1"
+DEFAULT_PROMPT = "1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt"
+
+def get_samples():
+    prompt_list = prompt_df.sample(n = 10)["Prompt"].map(lambda x: x).values.tolist()
+    return prompt_list
+
+def get_params(request: gr.Request):
+    params = request.query_params
+    ip = request.client.host
+    req = {"params": params,
+          "ip": ip}
+    return get_samples()
+
 list_models = [
     "SDXL-1.0",
     "SD-1.5",
@@ -39,6 +58,8 @@ def generate_txt2img(current_model, prompt, is_negative=False, image_style="None
     API_TOKEN = os.environ.get("HF_READ_TOKEN")
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
 
+    if type(prompt) != type(""):
+        prompt = DEFAULT_PROMPT
 
     if image_style == "None style":
         payload = {
@@ -181,7 +202,7 @@ with gr.Blocks(css=css) as demo:
     
     favicon = '<img src="" width="48px" style="display: inline">'
     gr.Markdown(
-        f"""<h1><center>{favicon} AI Diffusion</center></h1>
+        f"""<h1><center>🌻{favicon} AI Diffusion</center></h1>
             """
     )
     
@@ -189,8 +210,15 @@ with gr.Blocks(css=css) as demo:
         current_model = gr.Dropdown(label="Current Model", choices=list_models, value=list_models[1])
         
     with gr.Row(elem_id="prompt-container"):
-        text_prompt = gr.Textbox(label="Prompt", placeholder="a cute dog", lines=1, elem_id="prompt-text-input")
+        text_prompt = gr.Textbox(label="Prompt", placeholder="a cute dog", lines=1, elem_id="prompt-text-input", value = DEFAULT_PROMPT)
         text_button = gr.Button("Generate", variant='primary', elem_id="gen-button")
+
+    with gr.Row("prompt-container"):
+        select_prompt = gr.Dropdown(label="Prompt selected", choices=list_prompts, 
+                                   value = "1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt",
+                                   info = "default prompt: {}".format(DEFAULT_PROMPT)
+                                   )
+        select_button = gr.Button("Select Prompt Generate", variant='primary', elem_id="gen-button")
         
     with gr.Row():
         image_output = gr.Image(type="pil", label="Output Image", elem_id="gallery")
@@ -200,5 +228,7 @@ with gr.Blocks(css=css) as demo:
         image_style = gr.Dropdown(label="Style", choices=["None style", "Cinematic", "Digital Art", "Portrait"], value="None style", allow_custom_value=False)
 
     text_button.click(generate_txt2img, inputs=[current_model, text_prompt, negative_prompt, image_style], outputs=image_output)
+    select_button.click(generate_txt2img, inputs=[current_model, select_prompt, negative_prompt, image_style], outputs=image_output)
+    demo.load(get_params, None, select_prompt)
 
 demo.launch(show_api=False)