yamildiego
commited on
Commit
•
9e8370c
1
Parent(s):
2c11fb9
minor cghange
Browse files- handler.py +6 -6
handler.py
CHANGED
@@ -29,7 +29,7 @@ class EndpointHandler():
|
|
29 |
# import torch
|
30 |
|
31 |
device = "cuda"
|
32 |
-
num_images_per_prompt =
|
33 |
|
34 |
prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to(device)
|
35 |
decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to(device)
|
@@ -39,10 +39,10 @@ class EndpointHandler():
|
|
39 |
|
40 |
prior_output = prior(
|
41 |
prompt=prompt,
|
42 |
-
height=
|
43 |
-
width=
|
44 |
negative_prompt=negative_prompt,
|
45 |
-
guidance_scale=
|
46 |
num_images_per_prompt=num_images_per_prompt,
|
47 |
num_inference_steps=20
|
48 |
)
|
@@ -50,11 +50,11 @@ class EndpointHandler():
|
|
50 |
image_embeddings=prior_output.image_embeddings.half(),
|
51 |
prompt=prompt,
|
52 |
negative_prompt=negative_prompt,
|
53 |
-
guidance_scale=
|
54 |
output_type="pil",
|
55 |
num_inference_steps=10
|
56 |
).images
|
57 |
-
return decoder_output
|
58 |
|
59 |
|
60 |
|
|
|
29 |
# import torch
|
30 |
|
31 |
device = "cuda"
|
32 |
+
num_images_per_prompt = 1
|
33 |
|
34 |
prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to(device)
|
35 |
decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to(device)
|
|
|
39 |
|
40 |
prior_output = prior(
|
41 |
prompt=prompt,
|
42 |
+
height=512,
|
43 |
+
width=512,
|
44 |
negative_prompt=negative_prompt,
|
45 |
+
guidance_scale=7.0,
|
46 |
num_images_per_prompt=num_images_per_prompt,
|
47 |
num_inference_steps=20
|
48 |
)
|
|
|
50 |
image_embeddings=prior_output.image_embeddings.half(),
|
51 |
prompt=prompt,
|
52 |
negative_prompt=negative_prompt,
|
53 |
+
guidance_scale=7.0,
|
54 |
output_type="pil",
|
55 |
num_inference_steps=10
|
56 |
).images
|
57 |
+
return decoder_output[0]
|
58 |
|
59 |
|
60 |
|