app.py

#!/usr/bin/env python

import os  
import random  
import uuid  
import json

import gradio as gr  
import numpy as np  
from PIL import Image  
import spaces  
import torch  
from diffusers import DiffusionPipeline  
from typing import Tuple

# Check for the Model Base..

bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
default_negative = os.getenv("default_negative","")

def check_text(prompt, negative=""):
    for i in bad_words:
        if i in prompt:
            return True  
    for i in bad_words_negative:
        if i in negative:
            return True  
    return False

style_list = [
    {
        "name": "D&D Art",
        "prompt": "dungeons & dragons style artwork {prompt}. d&d style, key visual, vibrant, studio anime, highly detailed",
        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast"
    },
    {
        "name": "2560 x 1440",
        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic, fantastical environments, epic scenery, rich textures",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, sci-fi, modern cityscapes, deformed, ugly"
    },
    {
        "name": "Photo",
        "prompt": "cinematic photo of {prompt}. 35mm photograph, fantasy style, atmospheric lighting, magical elements, bokeh, professional, 4k, highly detailed",
        "negative_prompt": "modern day, sci-fi, minimalistic, noisy, blurry, overly dark, soft, deformed, low fantasy"
    },
    {
        "name": "Epic Battle",
        "prompt": "dynamic dungeons & dragons artwork of {prompt}. epic battle scene, legendary heroes, fierce monsters, intense action, dramatic lighting, high-detail environment, magical effects, vibrant colors",
        "negative_prompt": "peaceful, mundane, low energy, modern, sci-fi, simplistic, cartoonish, low contrast"
    },
    {
        "name": "Character Portrait",
        "prompt": "dungeons & dragons character portrait of {prompt}. intricate details, expressive faces, heroic poses, rich textures, elaborate costumes, iconic weapons, fantasy aesthetic, studio-quality rendering",
        "negative_prompt": "generic, overly stylized, blurry, simplistic, futuristic, deformed"
    },
    {
        "name": "Magical Realm",
        "prompt": "mystical dungeons & dragons artwork of {prompt}. enchanted forests, glowing runes, floating islands, otherworldly landscapes, magical auras, vibrant fantasy colors, ultra-detailed",
        "negative_prompt": "urban, dull, realistic, futuristic, lifeless, overly muted, minimalistic"
    },
    {
        "name": "Dark Fantasy",
        "prompt": "dark and moody dungeons & dragons artwork of {prompt}. gothic ruins, shadowy figures, haunting atmospheres, grim villains, cursed relics, muted colors, intricate textures, sinister undertones",
        "negative_prompt": "bright, cheerful, cartoonish, lighthearted, futuristic, deformed"
    },
    {
        "name": "Legendary Creatures",
        "prompt": "dungeons & dragons artwork of legendary creatures {prompt}. awe-inspiring dragons, ferocious beasts, mythical monsters, intricate scales, detailed textures, dynamic poses, magical auras, cinematic quality",
        "negative_prompt": "simplistic, dull, mundane creatures, low-detail, modern animals, cartoony"
    },
    {
        "name": "Ancient Ruins",
        "prompt": "dungeons & dragons artwork of {prompt}. ancient temples, crumbling ruins, arcane symbols, epic scale, mystical atmosphere, grand designs, dramatic perspective, ultra-detailed",
        "negative_prompt": "modern structures, low-detail, dull, untextured, non-fantasy elements, minimalistic"
    },
    {
        "name": "Spellcasting Scene",
        "prompt": "dungeons & dragons artwork of {prompt}. dramatic spellcasting, glowing magic effects, dynamic poses, swirling energy, vibrant light contrasts, powerful sorcerers, rich fantasy aesthetic",
        "negative_prompt": "low energy, non-magical, mundane, modern day, generic designs, dull colors"
    },
    {
        "name": "Tavern Life",
        "prompt": "dungeons & dragons artwork of {prompt}. lively medieval tavern scene, detailed character interactions, warm firelight, wooden interiors, mugs of ale, bard performances, rich textures, inviting atmosphere",
        "negative_prompt": "cold, empty, futuristic, lifeless, overly simplistic, boring"
    },
    {
        "name": "D&D Quest Art",
        "prompt": "dungeons & dragons quest artwork of {prompt}. heroic adventurers, perilous journeys, hidden treasures, legendary quests, immersive fantasy landscapes, intricate details, atmospheric lighting, grand scale, narrative-driven composition",
        "negative_prompt": "modern, futuristic, low-detail, overly abstract, mundane settings, simplistic, sci-fi, cartoony, lifeless, uninspired"
    },
    {
        "name": "D&D Anime Art",
        "prompt": "dungeons & dragons anime-style artwork of {prompt}. stylized characters, vibrant colors, expressive designs, dynamic action poses, magical effects, intricate fantasy backgrounds, studio-quality anime detailing",
        "negative_prompt": "photo-realistic, overly gritty, mundane, futuristic, overly dark, simplistic, deformed"
    },
    {
        "name": "D&D Map",
        "prompt": "dungeons & dragons map artwork of {prompt}. intricate fantasy map, visible grid layout for gameplay guidance, ancient cartography style, labeled locations, landmarks, topographic textures, rich colors, immersive design",
        "negative_prompt": "modern maps, bland, plain, low detail, futuristic, chaotic, inaccurate grids, unlabeled, minimalist"
    },
    {
        "name": "3D Model",
        "prompt": "professional 3D model of {prompt}. stylized fantasy character, legendary creatures, intricate armor, magical weapons, dynamic poses, volumetric lighting, epic fantasy feel, octane render, highly detailed",
        "negative_prompt": "low poly, blurry, noisy, simplistic, overly realistic, mundane objects, low effort"
    },
    {
        "name": "(No style)",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "D&D Quest Art"

def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    if not negative:
        negative = ""
    return p.replace("{prompt}", positive), n + negative

DESCRIPTION = """## ChatDnD.net Dungeons & Dragons Fantasy Creator  
Forge your own epic adventures and legendary heroes. Generate custom artwork for characters, quests, and magical landscapes to bring your campaigns to life. Finetuned by <a href='https://chatdnd.net'>ChatDnD</a>. To support me, <a href='https://buymeacoffee.com/watchoutformike'>Buy me a coffee</a>
"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"

MAX_SEED = np.iinfo(np.int32).max  
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

NUM_IMAGES_PER_PROMPT = 1

if torch.cuda.is_available():
    pipe = DiffusionPipeline.from_pretrained(
        "SG161222/RealVisXL_V3.0_Turbo",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16"
    )
    pipe2 = DiffusionPipeline.from_pretrained(
        "SG161222/RealVisXL_V2.02_Turbo",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16"
    )
    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
        pipe2.enable_model_cpu_offload()
    else:
        pipe.to(device)
        pipe2.to(device)
        print("Loaded on Device!")

    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
        print("Model Compiled!")

def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style: str = DEFAULT_STYLE_NAME,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3,
    use_resolution_binning: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    if check_text(prompt, negative_prompt):
        raise ValueError("Prompt contains restricted words.")

    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
    generator = torch.Generator().manual_seed(42)  # Fixed seed for reproducibility

    if not use_negative_prompt:
        negative_prompt = ""  # type: ignore  
    negative_prompt += default_negative

    options = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": 25,
        "generator": generator,
        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,
        "use_resolution_binning": use_resolution_binning,
        "output_type": "pil",
    }

    images = pipe(**options).images + pipe2(**options).images

    image_paths = [save_image(img) for img in images]
    return image_paths

examples = [
    "A grand Dungeons & Dragons quest scene: A heroic party of adventurers standing at the entrance of a towering, ancient stone castle covered in vines, under a sky lit by the fading orange hues of a setting sun. The characters include a fierce elven warrior in ornate armor, a wise wizard with a glowing staff, a stealthy halfling rogue with a bow, and a mighty dwarf cleric with a shining warhammer. In the background, a dark forest looms, and mist begins to rise from the ground. The ground around the castle is scattered with enchanted symbols and ancient runes. The atmosphere is tense and filled with magic, adventure, and mystery."
]

css = '''
body {
    background-color: #231f20;
    color: #f5f5f5;
    font-family: 'Arial', sans-serif;
}

.gradio-container {
    max-width: 900px !important; /* Increase width for a spacious layout */
    background: linear-gradient(180deg, #1a2332 0%, #08141e 100%);
    border-radius: 15px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.8);
    color: #d4e6f9;
    font-family: 'Roboto', sans-serif;
    padding: 20px;
}

h1, h2, h3 {
    color: #f4a261;
    font-family: 'Cinzel', serif; /* Give a medieval vibe */
    text-align: center;
    letter-spacing: 2px;
}

.gradio-slider {
    background-color: #0f1d29;
    border: 1px solid #3a506b;
    color: #d4e6f9;
}

.gr-checkbox label {
    color: #d4e6f9;
}

.gr-radio label {
    color: #f4a261;
    font-family: 'Cinzel', serif;
}

.gradio-button {
    background: #f4a261;
    color: #2a2a40;
    border-radius: 8px;
    box-shadow: 0 2px 10px rgba(244, 162, 97, 0.4);
    font-weight: bold;
}

.gradio-button:hover {
    background: #e76f51;
}

.gradio-gallery .label {
    color: #f4a261;
    font-family: 'Cinzel', serif;
    font-size: 1.2rem;
}

.accordion-content {
    background: #08141e;
    border: 1px solid #3a506b;
    padding: 15px;
}

.gr-row, .gr-column, .gr-group {
    border-radius: 10px;
}
'''

with gr.Blocks(css=css) as demo:
    gr.Markdown(f"""
    # ⚔️ Dungeons & Dragons Fantasy Generator ⚔️  
    **Unleash Your Imagination!** Create heroes, maps, quests, and epic scenes to bring your campaigns to life. Tailored for adventurers seeking inspiration or Dungeon Masters constructing their next grand story. <br>
    [Visit Our Website](https://chatdnd.net) | [Support Us](https://buymeacoffee.com/watchoutformike)
    *Roll for initiative and dive in!*
    """)
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="🔮 Your Prompt",
                show_label=True,
                max_lines=1,
                placeholder="Describe your vision (e.g., 'Elven Archer battling a hydra')",
                container=False,
            )
            run_button = gr.Button("Generate 🎨")
        result = gr.Gallery(label="Your Results", columns=2, preview=True)
    with gr.Accordion("📜 Advanced Options", open=False):
        use_negative_prompt = gr.Checkbox(label="🔻 Use Negative Prompt", value=True, visible=True)
        negative_prompt = gr.Text(
            label="Negative Prompt",
            max_lines=2,
            placeholder="Describe what to avoid (e.g., blurry, modern style)",
            value="",
            visible=True,
        )
        with gr.Row():
            num_inference_steps = gr.Slider(
                label="⚙️ Steps",
                minimum=10,
                maximum=60,
                step=1,
                value=30,
            )
        with gr.Row():
            num_images_per_prompt = gr.Slider(
                label="🖼️ Images Per Prompt",
                minimum=1,
                maximum=4,
                step=1,
                value=2,
            )
        with gr.Row(visible=True):
            width = gr.Slider(
                label="Width (px)",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
            height = gr.Slider(
                label="Height (px)",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="🎨 Guidance Scale",
                minimum=1.0,
                maximum=20.0,
                step=0.5,
                value=7.5,
            )
    with gr.Row(visible=True):
        style_selection = gr.Radio(
            label="🎭 Select Your Style",
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            container=True,
        )
    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=result,
        fn=generate,
        cache_examples=CACHE_EXAMPLES,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
        ],
        outputs=result,
        api_name="run",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()