Update app.py

app.py

@@ -1,174 +1,109 @@
+# app.py
+
 import gradio as gr
 import numpy as np
 import random
-
-import spaces  # [uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 import torch
+from diffusers import DiffusionPipeline
+from tags import TAGS
 
+# ========== Model Setup ==========
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v140-sdxl"  # Replace with your desired model
-
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+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.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-@spaces.GPU  # [uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    positive_prefix,
-    negative_prefix,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
+# ========== Inference Function ==========
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
+          guidance_scale, num_inference_steps, *tag_selections, active_tab):
+
+    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}"
+
+    additional_negatives = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
+    full_negative_prompt = f"{additional_negatives}, {negative_prompt}"
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
-    # Combine prefixes with prompts
-    full_prompt = positive_prefix + " " + prompt
-    full_negative_prompt = negative_prefix + " " + negative_prompt
-
     image = pipe(
-        prompt=full_prompt,
+        prompt=final_prompt,
         negative_prompt=full_negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,
         height=height,
-        generator=generator,
+        generator=generator
     ).images[0]
 
-    return image, seed
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Textbox(
-                label="Prompt",
-                show_label=False,
-                max_lines=5,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Textbox(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=True,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                positive_prefix = gr.Textbox(
-                    label="Positive prompt prefix",
-                    max_lines=3,
-                    placeholder="e.g., score_9, score_8_up, score_7_up, source_anime, ",
-                    value="score_9, score_8_up, score_7_up, source_anime, ",
-                )
-
-                negative_prefix = gr.Textbox(
-                    label="Negative prompt prefix",
-                    max_lines=3,
-                    placeholder="e.g., worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome,blurry, lowres, watermark, ",
-                    value="bad quality, worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome,blurry, lowres, watermark, ",
-                )
-
-            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():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=5.0,
-                    maximum=7.0,
-                    step=0.1,
-                    value=6.0,  # Adjusted default
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=20,
-                    maximum=30,
-                    step=1,
-                    value=25,
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
+    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:
+
+    with gr.Row():
+        with gr.Column(elem_id="left-column"):
+            gr.Markdown("# Rainbow Media X")
+
+            result = gr.Image(label="Result", show_label=False)
+            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")
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                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)
+
+                with gr.Row():
+                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=7.0)
+                    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")
+
+        with gr.Column(elem_id="right-column"):
+            active_tab = gr.State("Prompt Input")
+
+            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_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)
+                    tag_tab.select(lambda: "Tag Selection", outputs=active_tab)
+
+    run_button.click(
         fn=infer,
         inputs=[
-            prompt,
-            negative_prompt,
-            positive_prefix,
-            negative_prefix,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
+            prompt, negative_prompt, seed, randomize_seed, width, height,
+            guidance_scale, num_inference_steps,
+            *tag_checkboxes,
+            active_tab
         ],
-        outputs=[result, seed],
+        outputs=[result, seed, prompt_info]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.queue().launch()