Update app.py

app.py

@@ -9,8 +9,8 @@ import torch
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
 from huggingface_hub import hf_hub_download
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
-from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from openai import OpenAI
+import config
 
 # 从环境变量中获取 API 密钥
 api_key = os.getenv("MISTRAL_API_KEY")
@@ -57,12 +57,18 @@ css="""
 }
 """
 @spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True), guidance_scale=3.5):
+def infer(prompt, quality_select, styles_Radio, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True), guidance_scale=3.5):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
-    
+    if quality_select:
+        prompt += ", masterpiece, best quality, very aesthetic, absurdres"
+    if styles_Radio:
+        for style_name in styles_Radio:
+            for style in config.style_list:
+                if style["name"] == style_name:
+                    prompt += style["prompt"].replace("{prompt}", "")
     image = pipe(
            prompt = prompt, 
            width = width,
@@ -155,62 +161,69 @@ def predict(message, history, additional_dropdown):
 with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                placeholder="Enter your prompt",
-                container=False
-            )
-            run_button = gr.Button("Run")
-            result = gr.Image(label="Result", show_label=False, interactive=False)
-            
-            with gr.Accordion("Advanced Settings", open=False):
-            
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0,
+            with gr.Tab("Generator"):
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    placeholder="Enter your prompt",
+                    max_lines = 12,
+                    container=False
                 )
+                run_button = gr.Button("Run")
+                result = gr.Image(label="Result", show_label=False, interactive=False)
                 
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Accordion("Advanced Settings", open=False):
                 
-                with gr.Row():
-                    
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=1024,
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
                     )
                     
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=1024,
-                    )
-                
-                with gr.Row():
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     
-      
-                    num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=50,
-                        step=1,
-                        value=4,
-                    )
-                    guidancescale = gr.Slider(
-                        label="Guidance scale",
-                        minimum=0,
-                        maximum=10,
-                        step=0.1,
-                        value=3.5,
-                    )
+                    with gr.Row():
+                        
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1024,
+                        )
+                        
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1024,
+                        )
+                    
+                    with gr.Row():
+                        
+          
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=4,
+                        )
+                        guidancescale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=0,
+                            maximum=10,
+                            step=0.1,
+                            value=3.5,
+                        )
+            with gr.Tab("Styles"):
+                quality_select = gr.Checkbox(label="high quality")
+                styles_name = [style["name"] for style in config.style_list]
+                styles_Radio = gr.Dropdown(styles_name,label="Styles",multiselect=True)
+                
         with gr.Column(scale=3,elem_id="col-container"):
             gr.ChatInterface(
                 predict,
@@ -231,7 +244,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps, guidancescale],
+        inputs = [prompt, quality_select, styles_Radio, seed, randomize_seed, width, height, num_inference_steps, guidancescale],
         outputs = [result, seed]
     )
 if __name__ == "__main__":