Update app.py

app.py

@@ -15,10 +15,7 @@ DEFAULT_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 
 def load_pipeline(model):
     return DiffusionPipeline.from_pretrained(
-        model,
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant="fp16"
+        model, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
     )
 
 
@@ -29,12 +26,17 @@ loaded_models = {DEFAULT_MODEL}
 def generate_image(
     model: str,
     prompt: str,
+    negative_prompt: str,
+    num_inference_steps: int,
+    guidance_scale: float,
     progress,
 ):
     logger.debug(f"Loading pipeline: {dict(model=model)}")
     pipe = load_pipeline(model).to("cuda")
 
-    logger.debug(f"Generating image: {dict(prompt=prompt)}")
+    logger.debug(
+        f"Generating image: {dict(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)}"
+    )
 
     def progress_callback(pipe, step_index, timestep, callback_kwargs):
         logger.trace(
@@ -45,6 +47,9 @@ def generate_image(
 
     images = pipe(
         prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
         callback_on_step_end=progress_callback,
     ).images
     return images[0]
@@ -59,10 +64,13 @@ def gpu(*args, **kwargs):
 def generate(
     model: str,
     prompt: str,
+    negative_prompt: str,
+    num_inference_steps: int,
+    guidance_scale: float,
     progress=gr.Progress(),
 ):
     logger.info(
-        f"Starting image generation: {dict(model=model, prompt=prompt)}"
+        f"Starting image generation: {dict(model=model, prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)}"
     )
 
     # Cache the model files for the pipeline
@@ -72,7 +80,9 @@ def generate(
         load_pipeline(model)
         loaded_models.add(model)
 
-    return gpu(model, prompt, progress)
+    return gpu(
+        model, prompt, negative_prompt, num_inference_steps, guidance_scale, progress
+    )
 
 
 demo = gr.Interface(
@@ -82,8 +92,11 @@ demo = gr.Interface(
             label="Model", choices=models, value=DEFAULT_MODEL, allow_custom_value=True
         ),
         gr.Text(label="Prompt"),
+        gr.Text(label="Ngeative Prompt"),
+        gr.Slider(label="Num Inference Steps", minimum=0, maximum=200, value=50),
+        gr.Slider(label="Guidance Scale", minimum=0.0, maximum=300.0, value=7.5),
     ],
     outputs=[gr.Image(label="Output")],
 )
 
-demo.launch()
+demo.launch()