Update app.py

app.py

@@ -33,63 +33,81 @@ trigger_word = "Super Realism"  # Leave blank if not used
 pipe.load_lora_weights(lora_repo)
 pipe.to("cuda")
 
-# Define style options
+# Define style options with negative prompts
 style_list = [
     {
         "name": "3840 x 2160",
         "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "",
     },
     {
         "name": "2560 x 1440",
         "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "",
     },
     {
         "name": "HD+",
         "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "",
     },
     {
         "name": "Style Zero",
         "prompt": "{prompt}",
+        "negative_prompt": "",
     },
 ]
 
-styles = {k["name"]: k["prompt"] for k in style_list}
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 DEFAULT_STYLE_NAME = "3840 x 2160"
 STYLE_NAMES = list(styles.keys())
 
 # Apply selected style to the prompt
-def apply_style(style_name: str, positive: str) -> str:
-    return styles.get(style_name, styles[DEFAULT_STYLE_NAME]).replace("{prompt}", positive)
+def apply_style(style_name: str, positive: str) -> Tuple[str, str]:
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive), n
 
 # Image generation function with Spaces GPU support
 @spaces.GPU(duration=60, enable_queue=True)
 def generate(
     prompt: str,
+    negative_prompt: str = "",
+    use_negative_prompt: bool = False,
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
     guidance_scale: float = 3,
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
+    num_inference_steps: int = 30,
     progress=gr.Progress(track_tqdm=True),
 ):
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    positive_prompt = apply_style(style_name, prompt)
+    positive_prompt, style_negative_prompt = apply_style(style_name, prompt)
+    
+    if use_negative_prompt:
+        final_negative_prompt = style_negative_prompt + " " + negative_prompt
+    else:
+        final_negative_prompt = style_negative_prompt
+    
+    final_negative_prompt = final_negative_prompt.strip()
     
     if trigger_word:
         positive_prompt = f"{trigger_word} {positive_prompt}"
-
+    
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
     images = pipe(
         prompt=positive_prompt,
+        negative_prompt=final_negative_prompt if final_negative_prompt else None,
         width=width,
         height=height,
         guidance_scale=guidance_scale,
-        num_inference_steps=30,
+        num_inference_steps=num_inference_steps,
         num_images_per_prompt=1,
+        generator=generator,
         output_type="pil",
     ).images
     image_paths = [save_image(img) for img in images]
-    print(image_paths)
     return image_paths, seed
 
 # Example prompts
@@ -122,81 +140,87 @@ footer {
 '''
 
 # Gradio interface
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+with gr.Blocks(css=css) as demo:
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
+        prompt = gr.Text(
+            label="Prompt",
+            show_label=False,
+            max_lines=1,
+            placeholder="Enter your prompt",
+            container=False,
+        )
+        run_button = gr.Button("Run", scale=0, elem_classes="submit-btn")
+    result = gr.Gallery(label="Result", columns=1, show_label=False)
+    
+    with gr.Accordion("Advanced options", open=False):
+        style_selection = gr.Dropdown(
+            label="Quality Style",
+            choices=STYLE_NAMES,
+            value=DEFAULT_STYLE_NAME,
+            interactive=True,
+        )
+        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
+        negative_prompt = gr.Text(
+            label="Negative prompt",
+            max_lines=1,
+            placeholder="Enter a negative prompt",
+            visible=False,
+        )
+        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():
+            width = gr.Slider(
+                label="Width",
+                minimum=512,
+                maximum=2048,
+                step=64,
+                value=1280,
             )
-            run_button = gr.Button("Generated as 3:1 Images", scale=0, elem_classes="submit-btn")
-            
-            with gr.Accordion("Advanced options", open=True, visible=True):
-                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=64,
-                        value=1280,
-                    )
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=512,
-                        maximum=2048,
-                        step=64,
-                        value=832,
-                    )
-                
-                with gr.Row():
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=0.1,
-                        maximum=20.0,
-                        step=0.1,
-                        value=3.0,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=40,
-                        step=1,
-                        value=30,
-                    )
-
-                style_selection = gr.Radio(
-                    show_label=True,
-                    container=True,
-                    interactive=True,
-                    choices=STYLE_NAMES,
-                    value=DEFAULT_STYLE_NAME,
-                    label="Quality Style",
-                )
-        
-        with gr.Column(scale=2):
-            result = gr.Gallery(label="Result", columns=1, show_label=False)
-            
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                outputs=[result, seed],
-                fn=generate,
-                cache_examples=False,
+            height = gr.Slider(
+                label="Height",
+                minimum=512,
+                maximum=2048,
+                step=64,
+                value=832,
             )
-
+        guidance_scale = gr.Slider(
+            label="Guidance Scale",
+            minimum=0.1,
+            maximum=20.0,
+            step=0.1,
+            value=3.0,
+        )
+        num_inference_steps = gr.Slider(
+            label="Number of inference steps",
+            minimum=1,
+            maximum=40,
+            step=1,
+            value=30,
+        )
+    
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=[result, seed],
+        fn=generate,
+        cache_examples=False,
+    )
+    
+    # Handle visibility of negative prompt
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+        api_name=False,
+    )
+    
+    # Trigger generate on prompt submit or run button click
     gr.on(
         triggers=[
             prompt.submit,
@@ -205,12 +229,15 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         fn=generate,
         inputs=[
             prompt,
+            negative_prompt,
+            use_negative_prompt,
             seed,
             width,
             height,
             guidance_scale,
             randomize_seed,
             style_selection,
+            num_inference_steps,
         ],
         outputs=[result, seed],
         api_name="run",