WatchOutForMike/DnDQUEST2
0
1#!/usr/bin/env python2 3import os4import random5import uuid6import json7 8import gradio as gr9import numpy as np10from PIL import Image11import spaces12import torch13from diffusers import DiffusionPipeline14from typing import Tuple15 16#Check for the Model Base..//17 18 19 20bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))21bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))22default_negative = os.getenv("default_negative","")23 24def check_text(prompt, negative=""):25 for i in bad_words:26 if i in prompt:27 return True28 for i in bad_words_negative:29 if i in negative:30 return True31 return False32 33 34 35style_list = [36 37{38 "name": "D&D Art",39 "prompt": "dungeons & dragons style artwork {prompt}. d&d style, key visual, vibrant, studio anime, highly detailed",40 "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast"41},42{43 "name": "2560 x 1440",44 "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic, fantastical environments, epic scenery, rich textures",45 "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, sci-fi, modern cityscapes, deformed, ugly"46},47 48{49 "name": "Photo",50 "prompt": "cinematic photo of {prompt}. 35mm photograph, fantasy style, atmospheric lighting, magical elements, bokeh, professional, 4k, highly detailed",51 "negative_prompt": "modern day, sci-fi, minimalistic, noisy, blurry, overly dark, soft, deformed, low fantasy"52},53 54{55 "name": "Epic Battle",56 "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",57 "negative_prompt": "peaceful, mundane, low energy, modern, sci-fi, simplistic, cartoonish, low contrast"58},59 60{61 "name": "Character Portrait",62 "prompt": "dungeons & dragons character portrait of {prompt}. intricate details, expressive faces, heroic poses, rich textures, elaborate costumes, iconic weapons, fantasy aesthetic, studio-quality rendering",63 "negative_prompt": "generic, overly stylized, blurry, simplistic, futuristic, deformed"64},65 66{67 "name": "Magical Realm",68 "prompt": "mystical dungeons & dragons artwork of {prompt}. enchanted forests, glowing runes, floating islands, otherworldly landscapes, magical auras, vibrant fantasy colors, ultra-detailed",69 "negative_prompt": "urban, dull, realistic, futuristic, lifeless, overly muted, minimalistic"70},71 72{73 "name": "Dark Fantasy",74 "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",75 "negative_prompt": "bright, cheerful, cartoonish, lighthearted, futuristic, deformed"76},77 78{79 "name": "Legendary Creatures",80 "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",81 "negative_prompt": "simplistic, dull, mundane creatures, low-detail, modern animals, cartoony"82},83 84{85 "name": "Ancient Ruins",86 "prompt": "dungeons & dragons artwork of {prompt}. ancient temples, crumbling ruins, arcane symbols, epic scale, mystical atmosphere, grand designs, dramatic perspective, ultra-detailed",87 "negative_prompt": "modern structures, low-detail, dull, untextured, non-fantasy elements, minimalistic"88},89 90{91 "name": "Spellcasting Scene",92 "prompt": "dungeons & dragons artwork of {prompt}. dramatic spellcasting, glowing magic effects, dynamic poses, swirling energy, vibrant light contrasts, powerful sorcerers, rich fantasy aesthetic",93 "negative_prompt": "low energy, non-magical, mundane, modern day, generic designs, dull colors"94},95 96{97 "name": "Tavern Life",98 "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",99 "negative_prompt": "cold, empty, futuristic, lifeless, overly simplistic, boring"100},101 102{103 "name": "D&D Quest Art",104 "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",105 "negative_prompt": "modern, futuristic, low-detail, overly abstract, mundane settings, simplistic, sci-fi, cartoony, lifeless, uninspired"106},107 108{109 "name": "D&D Anime Art",110 "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",111 "negative_prompt": "photo-realistic, overly gritty, mundane, futuristic, overly dark, simplistic, deformed"112},113 114{115 "name": "D&D Map",116 "prompt": "dungeons & dragons map artwork of {prompt}. intricate fantasy map, visable grid layout for gameplay guidance, ancient cartography style, labeled locations, landmarks, topographic textures, rich colors, immersive design",117 "negative_prompt": "modern maps, bland, plain, low detail, futuristic, chaotic, inaccurate grids, unlabeled, minimalist"118},119 120{121 "name": "3D Model",122 "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",123 "negative_prompt": "low poly, blurry, noisy, simplistic, overly realistic, mundane objects, low effort"124},125 126 {127 "name": "(No style)",128 "prompt": "{prompt}",129 "negative_prompt": "",130 },131]132 133styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}134STYLE_NAMES = list(styles.keys())135DEFAULT_STYLE_NAME = "D&D Quest Art"136 137def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:138 p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])139 if not negative:140 negative = ""141 return p.replace("{prompt}", positive), n + negative142 143 144 145 146 147DESCRIPTION = """## ChatDnD.net Dungeons & Dragons Fantasy Creator 148 149Forge 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>150"""151 152 153 154 155 156if not torch.cuda.is_available():157 DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"158 159MAX_SEED = np.iinfo(np.int32).max160CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"161MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))162USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"163ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"164 165device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")166 167NUM_IMAGES_PER_PROMPT = 1168 169if torch.cuda.is_available():170 pipe = DiffusionPipeline.from_pretrained(171 "SG161222/RealVisXL_V3.0_Turbo",172 torch_dtype=torch.float16,173 use_safetensors=True,174 add_watermarker=False,175 variant="fp16"176 )177 pipe2 = DiffusionPipeline.from_pretrained(178 "SG161222/RealVisXL_V2.02_Turbo",179 torch_dtype=torch.float16,180 use_safetensors=True,181 add_watermarker=False,182 variant="fp16"183 )184 if ENABLE_CPU_OFFLOAD:185 pipe.enable_model_cpu_offload()186 pipe2.enable_model_cpu_offload()187 else:188 pipe.to(device) 189 pipe2.to(device) 190 print("Loaded on Device!")191 192 if USE_TORCH_COMPILE:193 pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)194 pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)195 print("Model Compiled!")196 197def save_image(img):198 unique_name = str(uuid.uuid4()) + ".png"199 img.save(unique_name)200 return unique_name201 202def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:203 if randomize_seed:204 seed = random.randint(0, MAX_SEED)205 return seed206 207@spaces.GPU(enable_queue=True)208def generate(209 prompt: str,210 negative_prompt: str = "",211 use_negative_prompt: bool = False,212 style: str = DEFAULT_STYLE_NAME,213 seed: int = 0,214 width: int = 1024,215 height: int = 1024,216 guidance_scale: float = 3,217 randomize_seed: bool = False,218 use_resolution_binning: bool = True,219 progress=gr.Progress(track_tqdm=True),220):221 if check_text(prompt, negative_prompt):222 raise ValueError("Prompt contains restricted words.")223 224 prompt, negative_prompt = apply_style(style, prompt, negative_prompt)225 seed = int(randomize_seed_fn(seed, randomize_seed))226 generator = torch.Generator().manual_seed(seed)227 228 if not use_negative_prompt:229 negative_prompt = "" # type: ignore230 negative_prompt += default_negative 231 232 options = {233 "prompt": prompt,234 "negative_prompt": negative_prompt,235 "width": width,236 "height": height,237 "guidance_scale": guidance_scale,238 "num_inference_steps": 25,239 "generator": generator,240 "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,241 "use_resolution_binning": use_resolution_binning,242 "output_type": "pil",243 }244 245 images = pipe(**options).images + pipe2(**options).images246 247 image_paths = [save_image(img) for img in images]248 return image_paths, seed249 250examples = [251 "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."252]253 254css = '''255.gradio-container {256 max-width: 900px !important; /* Increase width for a spacious layout */257 background: linear-gradient(180deg, #1a2332 0%, #08141e 100%);258 border-radius: 15px;259 box-shadow: 0 4px 15px rgba(0, 0, 0, 0.8);260 color: #d4e6f9;261 font-family: 'Roboto', sans-serif;262 padding: 20px;263}264 265h1, h2, h3 {266 color: #f4a261;267 font-family: 'Cinzel', serif; /* Give a medieval vibe */268 text-align: center;269 letter-spacing: 2px;270}271 272.gradio-slider {273 background-color: #0f1d29;274 border: 1px solid #3a506b;275 color: #d4e6f9;276}277 278.gr-checkbox label {279 color: #d4e6f9;280}281 282.gr-radio label {283 color: #f4a261;284 font-family: 'Cinzel', serif;285}286 287.gradio-button {288 background: #f4a261;289 color: #2a2a40;290 border-radius: 8px;291 box-shadow: 0 2px 10px rgba(244, 162, 97, 0.4);292 font-weight: bold;293}294 295.gradio-button:hover {296 background: #e76f51;297}298 299.gradio-gallery .label {300 color: #f4a261;301 font-family: 'Cinzel', serif;302 font-size: 1.2rem;303}304 305.accordion-content {306 background: #08141e;307 border: 1px solid #3a506b;308 padding: 15px;309}310 311.gr-row, .gr-column, .gr-group {312 border-radius: 10px;313}314'''315with gr.Blocks(css=css, theme="gradio_theme/ancient_look") as demo:316 gr.Markdown(f"""317 # ⚔️ Dungeons & Dragons Fantasy Generator ⚔️318 319 **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>320 [Visit Our Website](https://chatdnd.net) | [Support Us](https://buymeacoffee.com/watchoutformike)321 322 *Roll for initiative and dive in!*323 """)324 with gr.Group():325 with gr.Row():326 prompt = gr.Text(327 label="🔮 Your Prompt",328 show_label=True,329 max_lines=1,330 placeholder="Describe your vision (e.g., 'Elven Archer battling a hydra')",331 container=False,332 )333 run_button = gr.Button("Generate 🎨")334 result = gr.Gallery(label="Your Results", columns=2, preview=True)335 with gr.Accordion("📜 Advanced Options", open=False):336 use_negative_prompt = gr.Checkbox(label="🔻 Use Negative Prompt", value=True, visible=True)337 negative_prompt = gr.Text(338 label="Negative Prompt",339 max_lines=2,340 placeholder="Describe what to avoid (e.g., blurry, modern style)",341 value="",342 visible=True,343 )344 with gr.Row():345 num_inference_steps = gr.Slider(346 label="⚙️ Steps",347 minimum=10,348 maximum=60,349 step=1,350 value=30,351 )352 with gr.Row():353 num_images_per_prompt = gr.Slider(354 label="🖼️ Images Per Prompt",355 minimum=1,356 maximum=4,357 step=1,358 value=2,359 )360 seed = gr.Slider(361 label="🎲 Random Seed",362 minimum=0,363 maximum=10000,364 step=1,365 value=42,366 visible=True367 )368 randomize_seed = gr.Checkbox(label="🌟 Randomize Seed", value=True)369 with gr.Row(visible=True):370 width = gr.Slider(371 label="Width (px)",372 minimum=512,373 maximum=2048,374 step=8,375 value=1024,376 )377 height = gr.Slider(378 label="Height (px)",379 minimum=512,380 maximum=2048,381 step=8,382 value=1024,383 )384 with gr.Row():385 guidance_scale = gr.Slider(386 label="🎨 Guidance Scale",387 minimum=1.0,388 maximum=20.0,389 step=0.5,390 value=7.5,391 )392 with gr.Row(visible=True):393 style_selection = gr.Radio(394 label="🎭 Select Your Style",395 choices=STYLE_NAMES,396 value=DEFAULT_STYLE_NAME,397 container=True,398 )399 gr.Examples(400 examples=examples,401 inputs=prompt,402 outputs=[result, seed],403 fn=generate,404 cache_examples=CACHE_EXAMPLES,405 )406 407 use_negative_prompt.change(408 fn=lambda x: gr.update(visible=x),409 inputs=use_negative_prompt,410 outputs=negative_prompt,411 api_name=False,412 )413 414 gr.on(415 triggers=[416 prompt.submit,417 negative_prompt.submit,418 run_button.click,419 ],420 fn=generate,421 inputs=[422 prompt,423 negative_prompt,424 use_negative_prompt,425 style_selection,426 seed,427 width,428 height,429 guidance_scale,430 randomize_seed,431 ],432 outputs=[result, seed],433 api_name="run",434 )435 436if __name__ == "__main__":437 demo.queue(max_size=20).launch()438 439