Update app.py

app.py

@@ -3,15 +3,11 @@ import spaces
 import torch
 from diffusers import AuraFlowPipeline
 
-# Initialize the AuraFlow pipeline
 pipe = AuraFlowPipeline.from_pretrained("purplesmartai/pony-v7-base", torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
 @spaces.GPU()
 def generate_image(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed):
-    """
-    Function to generate an image using the AuraFlow pipeline.
-    """
     generator = torch.Generator("cuda").manual_seed(seed)
     image = pipe(
         prompt=prompt,
@@ -24,7 +20,6 @@ def generate_image(prompt, negative_prompt, height, width, num_inference_steps,
     ).images[0]
     return image
 
-# Create the Gradio interface
 iface = gr.Interface(
     fn=generate_image,
     inputs=[
@@ -36,10 +31,9 @@ iface = gr.Interface(
         gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.1, value=5.0),
         gr.Number(label="Seed", value=42)
     ],
-    outputs=gr.Image(label="Generated Image"),
-    title="AuraFlow Text-to-Image Generation",
+    outputs=gr.Image(label="Generated Image", format="png"),
+    title="purplesmartai/pony-v7-base",
     description="Generate images from text prompts using the AuraFlow model."
 )
 
-# Launch the Gradio app
 iface.launch()