Update app.py

app.py

@@ -39,8 +39,9 @@ def generate_image(prompt, ckpt):
         pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps==1 else "epsilon")
         pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda"))
         loaded = num_inference_steps
-        
-    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
+
+    timesteps = [399] if num_inference_steps == 1 else []
+    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=timesteps) 
 
 
     return results.images[0]