Update app.py

app.py

@@ -1,40 +1,40 @@
-# app.py
-
 import gradio as gr
 import numpy as np
 import random
 import torch
-import spaces  # Hugging Face Spaces ZeroGPU support
+import spaces
 from diffusers import DiffusionPipeline
-from tags import TAGS
+from tags import TAGS  # Centralized dictionary
 
-# Model loading
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
-model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-# ⬇ Inference Function with @spaces.GPU
+# Prepare keys for each tag category for UI and loop usage
+tag_categories = list(TAGS.keys())  # e.g. ["Participant", "Tribe", "Skin Tone", ...]
+tag_checkboxes = [gr.CheckboxGroup(choices=list(TAGS[k].keys()), label=f"{k} Tags") for k in tag_categories]
+
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
-          guidance_scale, num_inference_steps, *tag_selections, active_tab, progress=gr.Progress(track_tqdm=True)):
+          guidance_scale, num_inference_steps, active_tab, *tag_selections,
+          progress=gr.Progress(track_tqdm=True)):
 
     if active_tab == "Prompt Input":
         final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {prompt}"
     else:
-        all_tags = []
-        for (group_name, tag_dict), selected_keys in zip(TAGS.items(), tag_selections):
-            all_tags.extend([tag_dict[key] for key in selected_keys])
-        tag_text = ", ".join(all_tags)
-        final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {tag_text}"
+        combined_tags = []
+        for (tag_name, tag_dict), selected in zip(TAGS.items(), tag_selections):
+            combined_tags.extend([tag_dict[tag] for tag in selected])
+        tag_string = ", ".join(combined_tags)
+        final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {tag_string}"
 
-    additional_negatives = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
-    full_negative_prompt = f"{additional_negatives}, {negative_prompt}"
+    negative_base = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
+    full_negative_prompt = f"{negative_base}, {negative_prompt}"
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -53,13 +53,35 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
 
     return image, seed, f"Prompt used: {final_prompt}\nNegative prompt used: {full_negative_prompt}"
 
-# ⬇ Gradio UI
-with gr.Blocks(css="""
-#col-container { max-width: 1280px; margin: auto; }
-#left-column, #right-column { display: inline-block; vertical-align: top; width: 48%; padding: 1%; }
-#run-button { width: 100%; }
-""") as demo:
 
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1280px;
+}
+
+#left-column {
+    width: 50%;
+    display: inline-block;
+    padding: 20px;
+    vertical-align: top;
+}
+
+#right-column {
+    width: 50%;
+    display: inline-block;
+    vertical-align: top;
+    padding: 20px;
+    margin-top: 53px;
+}
+
+#run-button {
+    width: 100%;
+    margin-top: 10px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column(elem_id="left-column"):
             gr.Markdown("# Rainbow Media X")
@@ -68,16 +90,16 @@ with gr.Blocks(css="""
             prompt_info = gr.Textbox(label="Prompts Used", lines=3, interactive=False)
 
             with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt = gr.Textbox(label="Negative Prompt", max_lines=1, placeholder="Enter a negative prompt")
+                negative_prompt = gr.Textbox(label="Negative prompt", placeholder="Enter negative prompt")
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
                 with gr.Row():
-                    width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
-                    height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
+                    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():
-                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=7.0)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0, maximum=10, step=0.1, value=7)
                     num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=35)
 
             run_button = gr.Button("Run", elem_id="run-button")
@@ -87,26 +109,23 @@ with gr.Blocks(css="""
 
             with gr.Tabs() as tabs:
                 with gr.TabItem("Prompt Input") as prompt_tab:
-                    prompt = gr.Textbox(label="Prompt", placeholder="Enter a custom prompt", lines=3)
+                    prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt")
                     prompt_tab.select(lambda: "Prompt Input", outputs=active_tab)
 
                 with gr.TabItem("Tag Selection") as tag_tab:
-                    tag_checkboxes = []
-                    for group_name, tag_dict in TAGS.items():
-                        checkbox = gr.CheckboxGroup(choices=list(tag_dict.keys()), label=group_name)
-                        tag_checkboxes.append(checkbox)
+                    for tag_box in tag_checkboxes:
+                        tag_box.render()
                     tag_tab.select(lambda: "Tag Selection", outputs=active_tab)
 
-    run_button.click(
-        fn=infer,
-        inputs=[
-            prompt, negative_prompt, seed, randomize_seed,
-            width, height, guidance_scale, num_inference_steps,
-            active_tab,  # Moved BEFORE *tag_checkboxes
-            *tag_checkboxes
-        ],
-        outputs=[result, seed, prompt_info]
-    )
-
+        run_button.click(
+            fn=infer,
+            inputs=[
+                prompt, negative_prompt, seed, randomize_seed,
+                width, height, guidance_scale, num_inference_steps,
+                active_tab,
+                *tag_checkboxes
+            ],
+            outputs=[result, seed, prompt_info]
+        )
 
 demo.queue().launch()