Ziqi commited on
Commit
7ca7690
1 Parent(s): 4d0e27e
Files changed (1) hide show
  1. inference.py +5 -5
inference.py CHANGED
@@ -42,11 +42,11 @@ def make_image_grid(imgs, rows, cols):
42
 
43
 
44
  def inference_fn(
45
- model_id,
46
- prompt,
47
- num_samples,
48
- guidance_scale,
49
- ):
50
 
51
  # create inference pipeline
52
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
42
 
43
 
44
  def inference_fn(
45
+ model_id: str,
46
+ prompt: str,
47
+ num_samples: int,
48
+ guidance_scale: float,
49
+ ) -> PIL.Image.Image:
50
 
51
  # create inference pipeline
52
  device = 'cuda' if torch.cuda.is_available() else 'cpu'