Update app.py

app.py

@@ -40,10 +40,10 @@ def generate_image(prompt, ckpt):
         pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda"))
         loaded = num_inference_steps
 
-    if num_inference_steps == 1:
+    if loaded == 1:
         results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=[399])
     else:
-        results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
+        results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=[999, 749, 499, 249])
 
 
     return results.images[0]