update replicate

This commit is contained in:
Brett Kuprel 2022-07-12 12:04:56 -04:00
parent 51e401ea9b
commit 4e7b3b2caf

View File

@ -18,10 +18,11 @@ class ReplicatePredictor(BasePredictor):
def predict( def predict(
self, self,
text: str = Input(default='Dali painting of WALL·E'), text: str = Input(default='Dali painting of WALL·E'),
output_png: bool = Input(default=False),
intermediate_outputs: bool = Input(default=True), intermediate_outputs: bool = Input(default=True),
grid_size: int = Input(ge=1, le=9, default=5), grid_size: int = Input(ge=1, le=9, default=5),
log2_temperature: float = Input(ge=-3, le=3, default=0.0), log2_temperature: float = Input(ge=-3, le=3, default=2),
log2_top_k: int = Input(ge=0, le=14, default=7), log2_top_k: int = Input(ge=0, le=14, default=4),
log2_supercondition_factor: float = Input(ge=2, le=6, default=4) log2_supercondition_factor: float = Input(ge=2, le=6, default=4)
) -> Iterator[Path]: ) -> Iterator[Path]:
log2_mid_count = 3 if intermediate_outputs else 0 log2_mid_count = 3 if intermediate_outputs else 0
@ -40,7 +41,7 @@ class ReplicatePredictor(BasePredictor):
path = Path(tempfile.mkdtemp()) path = Path(tempfile.mkdtemp())
for image in image_stream: for image in image_stream:
i += 1 i += 1
ext = 'png' if i == 2 ** log2_mid_count else 'jpg' ext = 'png' if i == 2 ** log2_mid_count and output_png else 'jpg'
image_path = path / 'min-dalle-iter-{}.{}'.format(i, ext) image_path = path / 'min-dalle-iter-{}.{}'.format(i, ext)
image.save(str(image_path)) image.save(str(image_path))
yield image_path yield image_path