Update app.py

app.py

@@ -25,7 +25,20 @@ def select_model(model_name):
     if model_name in MODEL_LIST:
         return API_BASE_URL + model_name
 
-def query(prompt, selected_model, is_negative=False, steps=1, cfg_scale=6, seed=None):
+def extend_prompt(input_text):
+    API_URL = "https://api-inference.huggingface.co/models/Gustavosta/MagicPrompt-Stable-Diffusion"
+    
+    payload = {
+        "inputs": input_text
+    }
+
+    response = requests.post(API_URL, headers=HEADERS, json=payload)
+    response.raise_for_status()
+
+    extended_prompt = response.content
+    return extended_prompt
+
+def generate_image(prompt, selected_model, is_negative=False, steps=1, cfg_scale=6, seed=None):
     model_url = select_model(selected_model)
     API_URL = f"{model_url}"
 
@@ -63,7 +76,8 @@ with gr.Blocks(theme="soft") as playground:
         with gr.Column(elem_id="prompt-container"):
             text_prompt = gr.Textbox(label="Prompt", placeholder="a cute cat", lines=1, elem_id="prompt-text-input")
             model_dropdown = gr.Dropdown(label="Model", choices=MODEL_LIST, elem_id="model-dropdown", value="runwayml/stable-diffusion-v1-5")
-            text_button = gr.Button("Generate", variant='primary', elem_id="gen-button")
+            gen_button = gr.Button("Generate", variant='primary', elem_id="gen-button")
+            extend_button = gr.Button("Extend Prompt", variant='primary', elem_id="extend-button")
 
     with gr.Accordion("Advanced settings", open=False):
         negative_prompt = gr.Textbox(label="Negative Prompt", value="text, blurry, fuzziness", lines=1, elem_id="negative-prompt-text-input")
@@ -77,6 +91,7 @@ with gr.Blocks(theme="soft") as playground:
         """
     )
 
-    text_button.click(query, inputs=[text_prompt, model_dropdown, negative_prompt], outputs=image_output)
+    gen_button.click(generate_image, inputs=[text_prompt, model_dropdown, negative_prompt], outputs=image_output)
+    extend_button.click(extend_prompt, inputs=text_prompt, outputs=text_prompt)
 
 playground.launch(show_api=False)