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)