Update app.py

app.py

@@ -11,6 +11,7 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import DiffusionPipeline
+from typing import Tuple
 
 bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
 bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
@@ -25,6 +26,36 @@ def check_text(prompt, negative=""):
             return True
     return False
 
+style_list = [
+    {
+        "name": "3840 x 2160",
+        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "2560 x 1440",
+        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+
+    {
+        "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 = "3840 x 2160"
+
+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 = """"""
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
@@ -42,7 +73,6 @@ NUM_IMAGES_PER_PROMPT = 1
 if torch.cuda.is_available():
     pipe = DiffusionPipeline.from_pretrained(
         "SG161222/RealVisXL_V4.0",
-        ##"SG161222/RealVisXL_V1.02_Turbo",
         torch_dtype=torch.float16,
         use_safetensors=True,
         add_watermarker=False,
@@ -80,13 +110,12 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
-
 @spaces.GPU(enable_queue=True)
 def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
+    style: str = DEFAULT_STYLE_NAME,
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
@@ -95,7 +124,10 @@ def generate(
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipe.to(device)
+    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().manual_seed(seed)
 
@@ -104,32 +136,30 @@ def generate(
     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",
-
+        "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
+    images = pipe(**options).images + pipe2(**options).images
 
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
 
 examples = [
-    "neon holography crystal cat",
-    "a cat eating a piece of cheese",
-    "an astronaut riding a horse in space",
-    "a cartoon of a boy playing with a tiger",
-    "a cute robot artist painting on an easel, concept art",
-    #"a close up of a woman wearing a transparent, prismatic, elaborate nemeses headdress, over the should pose, brown skin-tone"
+    "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)",
+    "Silhouette of Hamburger standing in front of a, dark blue sky, a little saturated orange in the background sunset, night time, dark background, dark black hair, cinematic photography, cinematic lighting, dark theme, shattered camera lens, digital photography, 70mm, f2.8, lens aberration, grain, boke, double exposure, shaterred, color negative ",
+    "A silver grey cat wearing sunglasses and a tuxedo, as an agent in the movie The Matrix, in the style of a real photo --ar 43:31 --v 6.0 --style raw",
+    "A professional product photography of a (coffee cup, perfume bottle, hoodie) in front of a (vintage, modern) background --v 6.0 --style raw"
+
 ]
 
 css = '''
@@ -156,11 +186,12 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
         result = gr.Gallery(label="Result", columns=NUM_IMAGES_PER_PROMPT, show_label=False)
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
-            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=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, 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, NSFW",
                 visible=True,
             )
         seed = gr.Slider(
@@ -171,6 +202,7 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
             value=0,
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
         with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
@@ -221,6 +253,7 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
             prompt,
             negative_prompt,
             use_negative_prompt,
+            style,
             seed,
             width,
             height,
@@ -232,4 +265,4 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()