diff --git a/replicate_predictor.py b/replicate_predictor.py index 66e7d2f..e46be99 100644 --- a/replicate_predictor.py +++ b/replicate_predictor.py @@ -18,10 +18,11 @@ class ReplicatePredictor(BasePredictor): def predict( self, text: str = Input(default='Dali painting of WALL·E'), + output_png: bool = Input(default=False), intermediate_outputs: bool = Input(default=True), grid_size: int = Input(ge=1, le=9, default=5), - log2_temperature: float = Input(ge=-3, le=3, default=0.0), - log2_top_k: int = Input(ge=0, le=14, default=7), + log2_temperature: float = Input(ge=-3, le=3, default=2), + log2_top_k: int = Input(ge=0, le=14, default=4), log2_supercondition_factor: float = Input(ge=2, le=6, default=4) ) -> Iterator[Path]: log2_mid_count = 3 if intermediate_outputs else 0 @@ -40,7 +41,7 @@ class ReplicatePredictor(BasePredictor): path = Path(tempfile.mkdtemp()) for image in image_stream: 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.save(str(image_path)) yield image_path \ No newline at end of file