Update app.py

app.py

@@ -2,52 +2,115 @@ 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 diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from typing import Tuple
 
-css = '''
-.gradio-container {
-    max-width: 150%;
-    margin: 0 auto;
-}
-h1 { text-align: center; }
-footer { visibility: hidden; }
-'''
+bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
+bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
+default_negative = os.getenv("default_negative","")
 
-examples = [
-    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
-    "A glass cup of cold coffee placed on a rustic wooden table, surrounded by soft morning light. The coffee is rich, dark, and topped with a light layer of creamy froth, droplets of condensation sliding down the glass.",
-    "Vector illustration of a horse, vector graphic design with flat colors on an brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw",
-    "Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5",
-    "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 "
+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": "Photo",
+        "prompt": "cinematic photo {prompt}. 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
+    },   
+
+    {
+        "name": "Cinematic",
+        "prompt": "cinematic still {prompt}. emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+
+    {
+        "name": "Anime",
+        "prompt": "anime artwork {prompt}. anime style, key visual, vibrant, studio anime, highly detailed",
+        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
+    },
+    {
+        "name": "(No style)",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
 ]
 
-MODEL_ID = os.getenv("MODEL_VAL_PATH")  # SG161222/RealVisXL_V5.0_Lightning or SG161222/RealVisXL_V4.0_Lightning
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Photo"
+
+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 = """## Text to Image
+
+"""
+
+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"
-BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    use_safetensors=True,
-    add_watermarker=False,
-).to(device)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-if USE_TORCH_COMPILE:
-    pipe.compile()
+NUM_IMAGES_PER_PROMPT = 1
 
-if ENABLE_CPU_OFFLOAD:
-    pipe.enable_model_cpu_offload()
-
-MAX_SEED = np.iinfo(np.int32).max
+if torch.cuda.is_available():
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "SG161222/RealVisXL_V5.0_Lightning",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        add_watermarker=False,
+        variant="fp16"
+    )
+    pipe2 = StableDiffusionXLPipeline.from_pretrained(
+        "SG161222/RealVisXL_V4.0_Lightning",
+        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"
@@ -59,137 +122,168 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
-    seed: int = 1,
+    style: str = DEFAULT_STYLE_NAME,
+    seed: int = 0,
     width: int = 1024,
     height: int = 1024,
     guidance_scale: float = 3,
-    num_inference_steps: int = 25,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True,
-    num_images: int = 10,
     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)
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator().manual_seed(seed)
+
+    if not use_negative_prompt:
+        negative_prompt = ""  # type: ignore
+    negative_prompt += default_negative    
 
     options = {
-        "prompt": [prompt] * num_images,
-        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
+        "num_inference_steps": 25,
         "generator": generator,
+        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,
+        "use_resolution_binning": use_resolution_binning,
         "output_type": "pil",
     }
-
-    if use_resolution_binning:
-        options["use_resolution_binning"] = True
-
-    images = []
-    for i in range(0, num_images, BATCH_SIZE):
-        batch_options = options.copy()
-        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
-        if "negative_prompt" in batch_options:
-            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
-        images.extend(pipe(**batch_options).images)
+    
+    images = pipe(**options).images + pipe2(**options).images
 
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("## SDXL:Text to Image [10-Images]")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Row():
-                prompt = gr.Text(
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
-                run_button = gr.Button("Run", scale=0, variant="primary")
-
-            result = gr.Gallery(show_label=False, format="png", columns=2, object_fit="contain")
-
-            with gr.Accordion("Advanced Settings", open=False):
-                num_images = gr.Slider(
-                    label="Number of Images",
-                    minimum=1,
-                    maximum=15,
-                    step=1,
-                    value=10,
-                )
-                with gr.Row():
-                    use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-                    negative_prompt = gr.Text(
-                        label="Negative prompt",
-                        max_lines=5,
-                        lines=4,
-                        placeholder="Enter a negative prompt",
-                        value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-                        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(visible=True):
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=512,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=64,
-                        value=1024,
-                    )
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=512,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=64,
-                        value=1024,
-                    )
-                with gr.Row():
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=0.1,
-                        maximum=6,
-                        step=0.1,
-                        value=3.0,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=25,
-                        step=1,
-                        value=23,
-                    )
-
-        with gr.Column(scale=1):
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                cache_examples=False,
-            )
+examples = [
+    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
+    "A glass cup of cold coffee placed on a rustic wooden table, surrounded by soft morning light. The coffee is rich, dark, and topped with a light layer of creamy froth, droplets of condensation sliding down the glass.",
+    "Vector illustration of a horse, vector graphic design with flat colors on an brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw",
+    "Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5",
+    "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 "
+]
 
-        use_negative_prompt.change(
-            fn=lambda x: gr.update(visible=x),
-            inputs=use_negative_prompt,
-            outputs=negative_prompt,
-            api_name=False,
+css = '''
+.gradio-container {
+    max-width: 640px !important;
+    margin: 0 auto !important;
+}
+h1 {
+    text-align: center;
+}
+footer {
+    visibility: hidden;
+}
+'''
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Run")
+        result = gr.Gallery(label="Result", columns=1, 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=1,
+            placeholder="Enter a negative prompt",
+            value="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
+            visible=True,
         )
-    
+        with gr.Row():
+            num_inference_steps = gr.Slider(
+                label="Steps",
+                minimum=10,
+                maximum=60,
+                step=1,
+                value=20,
+            )
+        with gr.Row():
+            num_images_per_prompt = gr.Slider(
+                label="Images",
+                minimum=1,
+                maximum=10,
+                step=1,
+                value=6,
+            )
+        seed = gr.Slider(
+            label="Seed",
+            minimum=0,
+            maximum=MAX_SEED,
+            step=1,
+            value=0,
+            visible=True
+        )
+        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        with gr.Row(visible=True):
+            width = gr.Slider(
+                label="Width",
+                minimum=512,
+                maximum=2048,
+                step=8,
+                value=1024,
+            )
+            height = gr.Slider(
+                label="Height",
+                minimum=512,
+                maximum=2048,
+                step=8,
+                value=1024,
+            )
+        with gr.Row():
+            guidance_scale = gr.Slider(
+                label="Guidance Scale",
+                minimum=0.1,
+                maximum=20.0,
+                step=0.1,
+                value=3.0,
+            )
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value=DEFAULT_STYLE_NAME,
+            label="Image Style",
+        )
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=[result, seed],
+        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,
@@ -201,17 +295,16 @@ with gr.Blocks(css=css) as demo:
             prompt,
             negative_prompt,
             use_negative_prompt,
+            style_selection,
             seed,
             width,
             height,
             guidance_scale,
-            num_inference_steps,
             randomize_seed,
-            num_images
         ],
         outputs=[result, seed],
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=40).launch(ssr_mode=False)
+    demo.queue(max_size=20).launch(ssr_mode=True, show_error=True, share=True)