Update feifeiui/feifeiui.py

feifeiui/feifeiui.py

@@ -5,6 +5,7 @@ import config
 from feifeilib.feifeichat import feifeichat
 from feifeilib.feifeitexttoimg import feifeitexttoimg
 from feifeilib.feifeiflorence import feifeiflorence
+from feifeilib.feifeifluxapi import feifeifluxapi
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -21,7 +22,7 @@ def create_ui():
     with gr.Blocks(css=css) as FeiFei:
         with gr.Row():
             with gr.Column(scale=1):
-                with gr.Tab("Generator"):
+                with gr.Tab("FeiFei"):
                     prompt = gr.Text(
                         label="Prompt",
                         show_label=False,
@@ -29,7 +30,7 @@ def create_ui():
                         max_lines=12,
                         container=False,
                     )
-                    run_button = gr.Button("Run")
+                    feifei_button = gr.Button("FeiFei")
                     result = gr.Image(label="Result",
                                       show_label=False,
                                       interactive=False)
@@ -115,6 +116,70 @@ def create_ui():
                     output_img = gr.Image(label="Input Picture",
                                           interactive=False,
                                           show_label=False)
+                with gr.Tab("Flux"):
+                    flux_prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        placeholder="Enter your prompt",
+                        max_lines=12,
+                        container=False,
+                    )
+                    flux_button = gr.Button("Flux")
+                    flux_result = gr.Image(label="Result",
+                                      show_label=False,
+                                      interactive=False)
+
+                    with gr.Accordion("Advanced Settings", open=False):
+                        flux_seed = gr.Slider(
+                            label="Seed",
+                            minimum=0,
+                            maximum=MAX_SEED,
+                            step=1,
+                            value=0,
+                        )
+
+                        flux_randomize_seed = gr.Checkbox(label="Randomize seed",
+                                                     value=True)
+
+                        with gr.Row():
+                            flux_width = gr.Slider(
+                                label="Width",
+                                minimum=512,
+                                maximum=MAX_IMAGE_SIZE,
+                                step=64,
+                                value=1088,
+                            )
+
+                            flux_height = gr.Slider(
+                                label="Height",
+                                minimum=512,
+                                maximum=MAX_IMAGE_SIZE,
+                                step=64,
+                                value=1920,
+                            )
+
+                        with gr.Row():
+                            flux_num_inference_steps = gr.Slider(
+                                label="Number of inference steps",
+                                minimum=1,
+                                maximum=50,
+                                step=1,
+                                value=4,
+                            )
+                            flux_guidancescale = gr.Slider(
+                                label="Guidance scale",
+                                minimum=0,
+                                maximum=10,
+                                step=0.1,
+                                value=3.5,
+                            )
+                            flux_num_strength = gr.Slider(
+                                label="strength",
+                                minimum=0,
+                                maximum=2,
+                                step=0.001,
+                                value=0.035,
+                            )
             with gr.Column(scale=3, elem_id="col-container"):
                 gr.ChatInterface(
                     feifeichat,
@@ -141,7 +206,7 @@ def create_ui():
                     ],
                 )
 
-        run_button.click(
+        feifei_button.click(
             fn=feifeitexttoimg,  # Function to run for this button
             inputs=[
                 prompt,
@@ -159,7 +224,13 @@ def create_ui():
             ],
             outputs=[result, seed],
         )
-
+        flux_button.click(
+            fn=feifeifluxapi,  # Function to run for this button
+            inputs=[
+                flux_prompt
+            ],
+            outputs=[flux_result],
+        )
         florence_btn.click(
             fn=feifeiflorence,  # Function to run when the button is clicked
             inputs=[input_img],  # Input components for the function