Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusion3Pipeline
 import gradio as gr
 import os
 import spaces
@@ -14,19 +14,39 @@ else:
     print("Using CPU")
 
 # Initialize the pipeline and download the model
-pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN)
+pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN)
 pipe.to(device)
 
 # Define the image generation function
 @spaces.GPU(duration=60)
 def generate_image(prompt):
-    image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+    ).images[0]
     return image
 
 # Create the Gradio interface
+
+prompt = gr.Textbox(label="Prompt", info="Describe the image you want", placeholder="A cat...")
+
+negative_prompt = gr.Textbox(label="Negative Prompt", info="Describe what you don't want in the image", placeholder="Ugly, bad anatomy...")
+
+num_inference_steps = gr.Number(label="Number of Inference Steps", precision=0, value=25)
+
+height = gr.Number(label="Number of Inference Steps", precision=0, value=1024)
+
+width = gr.Number(label="Number of Inference Steps", precision=0, value=1024)
+
+guidance_scale = gr.Number(minimum=0.1, value=7.5, label="Guidance Scale", info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference")
+
 interface = gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(placeholder="Insert Prompt Here..."),
+    inputs=[prompt, negative_prompt, num_inference_steps, height, width, guidance_scale]
     outputs="image",
     title="Stable Diffusion 3 Medium",
     description="Made by [Nick088](https://linktr.ee/Nick088"