Update app.py

app.py

@@ -3,26 +3,35 @@ import spaces
 import torch
 from diffusers import AuraFlowPipeline
 import random
+import numpy as np
 
 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, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, sigmas_factor, seed, progress=gr.Progress(track_tqdm=True)):
     if seed < 0:
         seed = random.randint(0, 2**32 - 1)
 
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=int(height),
-        width=int(width),
-        num_inference_steps=int(num_inference_steps),
-        guidance_scale=guidance_scale,
-        generator=generator,
-    ).images[0]
+    pipeline_args = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "height": int(height),
+        "width": int(width),
+        "num_inference_steps": int(num_inference_steps),
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+    }
+
+    if sigmas_factor != 1.0:
+        steps = int(num_inference_steps)
+        sigmas = np.linspace(1.0, 1 / steps, steps)
+        sigmas = sigmas * sigmas_factor
+        pipeline_args["sigmas"] = sigmas.tolist()
+
+    image = pipe(**pipeline_args).images[0]
     return image, seed
 
 iface = gr.Interface(
@@ -34,6 +43,7 @@ iface = gr.Interface(
         gr.Slider(label="Width", minimum=256, maximum=2048, step=64, value=1024),
         gr.Slider(label="Number of Inference Steps", minimum=1, maximum=100, step=1, value=30),
         gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.1, value=3.5),
+        gr.Slider(label="Sigmas Factor", minimum=0.1, maximum=2.0, step=0.01, value=1.0),
         gr.Number(label="Seed (set to -1 for random)", value=-1, minimum=-1)
     ],
     outputs=[