Add Negative Prompt parameter

app.py

@@ -55,20 +55,20 @@ def calculate_resolution(aspect_ratio, mode='landscape', total_pixels=1024*1024,
 
 # Example prompts with ckpt, aspect, and mode
 examples = [
-    {"prompt": "A futuristic cityscape at sunset", "ckpt": "4-Step", "aspect": "16:9", "mode": "landscape"},
-    {"prompt": "pair of shoes made of dried fruit skins, 3d render, bright colours, clean composition, beautiful artwork, logo", "ckpt": "2-Step", "aspect": "1:1", "mode": "portrait"},
-    {"prompt": "A portrait of a robot in the style of Renaissance art", "ckpt": "2-Step", "aspect": "1:1", "mode": "portrait"},
-    {"prompt": "full body of alien shaped like woman, big golden eyes, mars planet, photo, digital art, fantasy", "ckpt": "4-Step", "aspect": "1:1", "mode": "portrait"},    
-    {"prompt": "A serene landscape with mountains and a river", "ckpt": "8-Step", "aspect": "3:2", "mode": "landscape"},
-    {"prompt": "post-apocalyptic wasteland, the most delicate beautiful flower with green leaves growing from dust and rubble, vibrant colours, cinematic", "ckpt": "8-Step", "aspect": "16:9", "mode": "landscape"}
+    {"prompt": "A futuristic cityscape at sunset", "negative_prompt": "Ugly", "ckpt": "4-Step", "aspect": "16:9", "mode": "landscape"},
+    {"prompt": "pair of shoes made of dried fruit skins, 3d render, bright colours, clean composition, beautiful artwork, logo", "negative_prompt": "Ugly", "ckpt": "2-Step", "aspect": "1:1", "mode": "portrait"},
+    {"prompt": "A portrait of a robot in the style of Renaissance art", "negative_prompt": "Ugly", "ckpt": "2-Step", "aspect": "1:1", "mode": "portrait"},
+    {"prompt": "full body of alien shaped like woman, big golden eyes, mars planet, photo, digital art, fantasy", "negative_prompt": "Ugly", "ckpt": "4-Step", "aspect": "1:1", "mode": "portrait"},    
+    {"prompt": "A serene landscape with mountains and a river", "negative_prompt": "Ugly", "ckpt": "8-Step", "aspect": "3:2", "mode": "landscape"},
+    {"prompt": "post-apocalyptic wasteland, the most delicate beautiful flower with green leaves growing from dust and rubble, vibrant colours, cinematic", "negative_prompt": "Ugly", "ckpt": "8-Step", "aspect": "16:9", "mode": "landscape"}
 ]
 # Define a function to set the example inputs
 def set_example(selected_prompt):
     # Find the example that matches the selected prompt
     for example in examples:
         if example["prompt"] == selected_prompt:
-            return example["prompt"], example["ckpt"], example["aspect"], example["mode"]
-    return None, None, None, None  # Default values if not found
+            return example["prompt"], example["negative_prompt"], example["ckpt"], example["aspect"], example["mode"]
+    return None, None, None, None, None  # Default values if not found
 
 # Check if CUDA is available (GPU support), and set the appropriate device
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -114,7 +114,7 @@ def check_nsfw_images(
 
 # Function 
 @spaces.GPU(enable_queue=True)
-def generate_image(prompt, ckpt, aspect_ratio, mode):
+def generate_image(prompt, negative_prompt, ckpt, aspect_ratio, mode):
     width, height = calculate_resolution(aspect_ratio, mode)  # Calculate resolution based on the aspect ratio
     checkpoint = checkpoints[ckpt][0]
     num_inference_steps = checkpoints[ckpt][1] 
@@ -127,7 +127,7 @@ def generate_image(prompt, ckpt, aspect_ratio, mode):
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
         
     pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device=device))
-    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, width=width, height=height )
+    results = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=0, width=width, height=height )
 
     if SAFETY_CHECKER:
         images, has_nsfw_concepts = check_nsfw_images(results.images)
@@ -152,6 +152,8 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Group():
         with gr.Row():
             prompt = gr.Textbox(label='Enter you image prompt:', scale=8)
+        with gr.Row():
+            negative_prompt = gr.Textbox(label='Optional negative prompt:', scale=8)
         with gr.Row():            
             ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
             aspect = gr.Dropdown(label='Aspect Ratio', choices=list(aspect_ratios.keys()), value='1:1', interactive=True)
@@ -161,15 +163,15 @@ with gr.Blocks(css="style.css") as demo:
     img = gr.Image(label='SDXL-Lightning Generated Image')
 
     prompt.submit(fn=generate_image,
-                 inputs=[prompt, ckpt, aspect, mode],
+                 inputs=[prompt, negative_prompt, ckpt, aspect, mode],
                  outputs=img,
                  )
     submit.click(fn=generate_image,
-                 inputs=[prompt, ckpt, aspect, mode],
+                 inputs=[prompt, negative_prompt, ckpt, aspect, mode],
                  outputs=img,
                  )
     # Dropdown for selecting examples
     example_dropdown = gr.Dropdown(label='Select an Example', choices=[e["prompt"] for e in examples])
-    example_dropdown.change(fn=set_example, inputs=example_dropdown, outputs=[prompt, ckpt, aspect, mode])
+    example_dropdown.change(fn=set_example, inputs=example_dropdown, outputs=[prompt, negative_prompt, ckpt, aspect, mode])
 
 demo.queue().launch()