Update app.py

app.py

@@ -9,6 +9,10 @@ import bs4
 import lxml
 # Define the list of models
 models = ["Qwen/Qwen2.5-Coder-32B-Instruct"]
+im_models = [
+    "black-forest-labs/FLUX.1-dev",
+    "ostris/Flex.1-alpha",
+]
 def format_prompt(message, mod, system):
     eos=f"<|im_end|>\n"
     bos=f"<|im_start|>\n"
@@ -58,7 +62,7 @@ def generate_prompt(company_name, company_html, company_descp):
     return output
 
         
-def generate_image(comp_name, comp_desc="", comp_url=""):
+def generate_image(comp_name, comp_desc="", comp_url="", im_model=""):
     print(comp_name +"|"+comp_desc+"|"+comp_url)
     if comp_url != "":
         web_out=get_webpage_text(comp_url)
@@ -66,7 +70,7 @@ def generate_image(comp_name, comp_desc="", comp_url=""):
         web_out="No Web Data returned"
     print(web_out)
     prompt_out=generate_prompt(comp_name,web_out, comp_desc)
-    client = InferenceClient("black-forest-labs/FLUX.1-dev")
+    client = InferenceClient(im_model)
     response = client.text_to_image(prompt_out)
     return response
 
@@ -78,6 +82,7 @@ with gr.Blocks() as demo:
             comp_name = gr.Textbox(label="Enter Company Name")
             comp_desc = gr.Textbox(label="Enter Company Description")
             comp_url = gr.Textbox(label="Enter Company URL")
+            im_model = gr.Dropdown(label="Image Model", choices=im_models)
             generate_button = gr.Button("Generate Image")
         with gr.Column():
             output_image = gr.Image(label="Generated Image")