emeersman commited on
Commit
0f9de12
1 Parent(s): 24ba3d1

Conditionally add random generator to inference pipeline

Browse files
Files changed (1) hide show
  1. handler.py +18 -7
handler.py CHANGED
@@ -36,13 +36,13 @@ class EndpointHandler():
36
  width = params.pop("width", None)
37
  manual_seed = params.pop("manual_seed", -1)
38
 
39
- generator = torch.Generator(device)
40
 
41
  if manual_seed != -1:
 
42
  generator.manual_seed(manual_seed)
43
-
44
- # run inference pipeline
45
- out = self.pipe(prompt,
46
  generator=generator,
47
  num_inference_steps=num_inference_steps,
48
  guidance_scale=guidance_scale,
@@ -50,7 +50,18 @@ class EndpointHandler():
50
  negative_prompt=negative_prompt,
51
  height=height,
52
  width=width
53
- )
54
-
55
- # return first generate PIL image
 
 
 
 
 
 
 
 
 
 
 
56
  return out.images[0]
 
36
  width = params.pop("width", None)
37
  manual_seed = params.pop("manual_seed", -1)
38
 
39
+ out = None
40
 
41
  if manual_seed != -1:
42
+ generator = torch.Generator(device='cuda')
43
  generator.manual_seed(manual_seed)
44
+ # run inference pipeline
45
+ out = self.pipe(prompt,
 
46
  generator=generator,
47
  num_inference_steps=num_inference_steps,
48
  guidance_scale=guidance_scale,
 
50
  negative_prompt=negative_prompt,
51
  height=height,
52
  width=width
53
+ )
54
+ else:
55
+ # run inference pipeline
56
+ out = self.pipe(prompt,
57
+ num_inference_steps=num_inference_steps,
58
+ guidance_scale=guidance_scale,
59
+ num_images_per_prompt=1,
60
+ negative_prompt=negative_prompt,
61
+ height=height,
62
+ width=width
63
+ )
64
+
65
+
66
+ # return first generated PIL image
67
  return out.images[0]