Yardenfren commited on
Commit
77e3547
1 Parent(s): 7c1777a

Update inf.py

Browse files
Files changed (1) hide show
  1. inf.py +3 -1
inf.py CHANGED
@@ -105,6 +105,7 @@ class InferencePipeline:
105
  seed: int,
106
  n_steps: int,
107
  guidance_scale: float,
 
108
  ) -> PIL.Image.Image:
109
  if not torch.cuda.is_available():
110
  raise gr.Error('CUDA is not available.')
@@ -117,5 +118,6 @@ class InferencePipeline:
117
  num_inference_steps=n_steps,
118
  guidance_scale=guidance_scale,
119
  generator=generator,
 
120
  ) # type: ignore
121
- return out.images[0]
 
105
  seed: int,
106
  n_steps: int,
107
  guidance_scale: float,
108
+ num_images_per_prompt: int = 1
109
  ) -> PIL.Image.Image:
110
  if not torch.cuda.is_available():
111
  raise gr.Error('CUDA is not available.')
 
118
  num_inference_steps=n_steps,
119
  guidance_scale=guidance_scale,
120
  generator=generator,
121
+ num_images_per_prompt=num_images_per_prompt,
122
  ) # type: ignore
123
+ return out.images