Update app.py

app.py

@@ -4,6 +4,7 @@ import random
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline, DPMSolverSDEScheduler
 import torch
+from tags import tag_options_1, tag_options_2  # Import tags here
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Replace to the model you would like to use
@@ -14,34 +15,14 @@ else:
     torch_dtype = torch.float32
 
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-# pipe.scheduler = DPMSolverSDEScheduler.from_config(pipe.scheduler.config, algorithm_type="dpmsolver++", solver_order=2, use_karras_sigmas=True)
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-# First set of tags
-tag_options_1 = {
-    "Fantasy": "fantasy",
-    "Sci-Fi": "sci-fi",
-    "Realistic": "realistic",
-    "Cyberpunk": "cyberpunk",
-    "Noir": "noir",
-}
-
-# Second set of tags
-tag_options_2 = {
-    "Vibrant Colors": "vibrant colors",
-    "Dark Theme": "dark theme",
-    "Minimalist": "minimalist",
-    "Photorealistic": "photorealistic",
-    "Abstract": "abstract",
-}
-
 @spaces.GPU  # [uncomment to use ZeroGPU]
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, tag_selection_1, tag_selection_2, use_tags, progress=gr.Progress(track_tqdm=True)):
 
-    # Determine prompt based on the tab selection
     if use_tags:
         selected_tags_1 = [tag_options_1[tag] for tag in tag_selection_1 if tag in tag_options_1]
         selected_tags_2 = [tag_options_2[tag] for tag in tag_selection_2 if tag in tag_options_2]
@@ -87,9 +68,11 @@ with gr.Blocks(css=css) as demo:
         # Text-to-Image Gradio Template
         """)
 
+        # Display result image at the top
+        result = gr.Image(label="Result", show_label=False)
+
         # Tabbed interface to select either Prompt or Tags
         with gr.Tabs() as tabs:
-            # Prompt-only tab
             with gr.TabItem("Prompt Input"):
                 prompt = gr.Text(
                     label="Prompt",
@@ -98,17 +81,18 @@ with gr.Blocks(css=css) as demo:
                     placeholder="Enter your prompt",
                     container=False,
                 )
-                use_tags = gr.State(False)  # This keeps track of whether tags are used
+                use_tags = gr.State(False)
 
-            # Tag-based input tab
             with gr.TabItem("Tag Selection"):
-                with gr.Row():
+                # Separate each tag section vertically
+                with gr.Column():
                     tag_selection_1 = gr.CheckboxGroup(choices=list(tag_options_1.keys()), label="Select Tags (Style)")
+                with gr.Column():
                     tag_selection_2 = gr.CheckboxGroup(choices=list(tag_options_2.keys()), label="Select Tags (Theme)")
-                use_tags = gr.State(True)  # When this tab is selected, tags are used
+                use_tags = gr.State(True)
 
-        run_button = gr.Button("Run", scale=0)
-        result = gr.Image(label="Result", show_label=False)
+        # Full-width "Run" button
+        run_button = gr.Button("Run", scale=0, full_width=True)
 
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
@@ -134,7 +118,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=1024,
                 )
 
                 height = gr.Slider(
@@ -142,7 +126,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=1024,
                 )
 
             with gr.Row():
@@ -151,7 +135,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=7,  # Replace with defaults that work for your model
+                    value=7,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -159,7 +143,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=35,  # Replace with defaults that work for your model
+                    value=35,
                 )
 
         gr.Examples(
@@ -167,13 +151,9 @@ with gr.Blocks(css=css) as demo:
             inputs=[prompt]
         )
 
-    # Use `run_button.click` event to check which tab is selected
     def check_tab(prompt, tag_selection_1, tag_selection_2, selected_tab):
-        if selected_tab == "Prompt Input":
-            return False  # Use prompt only
-        return True  # Use tags if on "Tag Selection"
+        return selected_tab == "Tag Selection"
 
-    # Set the selected tab
     tabs.change(check_tab, inputs=[prompt, tag_selection_1, tag_selection_2, tabs], outputs=use_tags)
 
     gr.on(