add model radio

myapp/app.py

@@ -19,6 +19,11 @@ except ImportError:
 
 client = InferenceClient(model="black-forest-labs/FLUX.1-schnell")
 
+MODELS = [
+    "stabilityai/stable-diffusion-3.5-large",
+    "black-forest-labs/FLUX.1-schnell",
+]
+
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
@@ -26,31 +31,17 @@ with gr.Blocks() as demo:
                 "https://wheelingvultures.bandcamp.com/album/ep", label="Text"
             )
             prompt = gr.TextArea("A psychedelic vulture", label="Prompt")
+            model = gr.Radio(MODELS, value=MODELS[0], label="Model")
             button = gr.Button("Generate")
         with gr.Column():
             background = gr.Image(visible=False, type="filepath")
             scale = gr.Slider(3, 15, 9, step=1, label="Scale")
             output = gr.Image()
 
-    gr.on(
-        [button.click, prompt.submit],
-        partial(gr.update, interactive=False),
-        outputs=button,
-    ).then(
-        partial(client.text_to_image, width=400, height=400),
-        inputs=prompt,
-        outputs=background,
-    ).then(
-        partial(gr.update, interactive=True),
-        outputs=button,
-    )
+    def generate_background(data: dict[Component, Any]):
+        return client.text_to_image(data[prompt], model=data[model]), None
 
-    @gr.on(
-        triggers=[text.submit, background.change, scale.change],
-        inputs={text, background, scale},
-        outputs=output,
-    )
-    def generate_code(data: dict[Component, Any]):
+    def generate_output(data: dict[Component, Any]):
         if data.get(background) is None:
             return None
 
@@ -67,6 +58,26 @@ with gr.Blocks() as demo:
 
             return Image.open(temp_file.name)
 
+    gr.on(
+        [button.click, prompt.submit],
+        partial(gr.update, interactive=False),
+        outputs=button,
+    ).then(
+        generate_background,
+        inputs={prompt, model},
+        outputs=[background, output],
+    ).then(
+        partial(gr.update, interactive=True),
+        outputs=button,
+    )
+
+    gr.on(
+        [text.submit, background.change, scale.change],
+        generate_output,
+        inputs={text, background, scale},
+        outputs=output,
+    )
+
 
 if __name__ == "__main__":
     demo.launch()