From 88676880fbb5a25f0d7ea34ebc2295a308307225 Mon Sep 17 00:00:00 2001 From: asutermo Date: Thu, 9 Jan 2025 16:32:48 -0700 Subject: [PATCH 1/6] support multi-image --- predict.py | 48 ++++++++++++++++++++++++++---------------------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/predict.py b/predict.py index 8ef6129..511d6a1 100644 --- a/predict.py +++ b/predict.py @@ -5,6 +5,7 @@ import subprocess import time import sys +from typing import List from cog import BasePredictor, Input, Path from PIL import Image @@ -96,31 +97,34 @@ def predict( description="Automatically adjust the output image size to be same as input image size. For editing and controlnet task, it can make sure the output image has the same size as input image leading to better performance", default=False, ), - ) -> Path: + num_images: int = Input(description="The number of images to generate for the given inputs", default=1, ge=1, le=10), + ) -> List[Path]: """Run a single prediction on the model""" if seed is None: seed = int.from_bytes(os.urandom(2), "big") print(f"Using seed: {seed}") input_images = [str(img) for img in [img1, img2, img3] if img is not None] - - output = self.pipe( - prompt=prompt, - input_images=None if len(input_images) == 0 else input_images, - height=height, - width=width, - guidance_scale=guidance_scale, - img_guidance_scale=img_guidance_scale, - num_inference_steps=inference_steps, - separate_cfg_infer=separate_cfg_infer, - use_kv_cache=True, - offload_kv_cache=True, - offload_model=offload_model, - use_input_image_size_as_output=use_input_image_size_as_output, - seed=seed, - max_input_image_size=max_input_image_size, - ) - img = output[0] - out_path = "/tmp/out.png" - img.save(out_path) - return Path(out_path) + output = [] + for i in range(num_images): + output = self.pipe( + prompt=prompt, + input_images=None if len(input_images) == 0 else input_images, + height=height, + width=width, + guidance_scale=guidance_scale, + img_guidance_scale=img_guidance_scale, + num_inference_steps=inference_steps, + separate_cfg_infer=separate_cfg_infer, + use_kv_cache=True, + offload_kv_cache=True, + offload_model=offload_model, + use_input_image_size_as_output=use_input_image_size_as_output, + seed=seed, + max_input_image_size=max_input_image_size, + ) + img = output[0] + out_path = f"/tmp/out_{i}.png" + img.save(out_path) + output.append(Path(out_path)) + return output From 700908ca6356dc1210d456832d5565d49b46beea Mon Sep 17 00:00:00 2001 From: asutermo Date: Thu, 9 Jan 2025 16:36:34 -0700 Subject: [PATCH 2/6] start working on updating pipeline --- OmniGen/pipeline.py | 1 + 1 file changed, 1 insertion(+) diff --git a/OmniGen/pipeline.py b/OmniGen/pipeline.py index 572452b..1a3067e 100644 --- a/OmniGen/pipeline.py +++ b/OmniGen/pipeline.py @@ -154,6 +154,7 @@ def __call__( dtype: torch.dtype = torch.bfloat16, seed: int = None, output_type: str = "pil", + num_images: int = 1, ): r""" Function invoked when calling the pipeline for generation. From 679e900af7d3af8fb5817f544051e2d5eca0cc0c Mon Sep 17 00:00:00 2001 From: asutermo Date: Thu, 9 Jan 2025 16:37:23 -0700 Subject: [PATCH 3/6] pipeline --- OmniGen/pipeline.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/OmniGen/pipeline.py b/OmniGen/pipeline.py index 1a3067e..4345c1d 100644 --- a/OmniGen/pipeline.py +++ b/OmniGen/pipeline.py @@ -193,6 +193,8 @@ def __call__( data type for the model output_type (`str`, *optional*, defaults to "pil"): The type of the output image, which can be "pt" or "pil" + num_images (`int`, *optional*, defaults to 1): + The number of images to generate Examples: Returns: From 8cab8de384c361aece94e48dec81b039a6a01cb6 Mon Sep 17 00:00:00 2001 From: asutermo Date: Thu, 9 Jan 2025 16:39:57 -0700 Subject: [PATCH 4/6] num_iter --- predict.py | 1 + 1 file changed, 1 insertion(+) diff --git a/predict.py b/predict.py index 511d6a1..921e8d4 100644 --- a/predict.py +++ b/predict.py @@ -122,6 +122,7 @@ def predict( use_input_image_size_as_output=use_input_image_size_as_output, seed=seed, max_input_image_size=max_input_image_size, + num_images=num_images, ) img = output[0] out_path = f"/tmp/out_{i}.png" From 38b6eaea16018b9566c9dd3a3ffd7f2f0cdb0775 Mon Sep 17 00:00:00 2001 From: asutermo Date: Fri, 10 Jan 2025 13:46:43 -0700 Subject: [PATCH 5/6] remove unecessary param, update gitignore --- .gitignore | 3 +++ OmniGen/pipeline.py | 3 --- predict.py | 7 +++---- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 79174c8..ab24571 100644 --- a/.gitignore +++ b/.gitignore @@ -172,3 +172,6 @@ local.ipynb convert_to_safetensor.py test_memory_time.py test.py + +# cog +.cog/ \ No newline at end of file diff --git a/OmniGen/pipeline.py b/OmniGen/pipeline.py index 4345c1d..572452b 100644 --- a/OmniGen/pipeline.py +++ b/OmniGen/pipeline.py @@ -154,7 +154,6 @@ def __call__( dtype: torch.dtype = torch.bfloat16, seed: int = None, output_type: str = "pil", - num_images: int = 1, ): r""" Function invoked when calling the pipeline for generation. @@ -193,8 +192,6 @@ def __call__( data type for the model output_type (`str`, *optional*, defaults to "pil"): The type of the output image, which can be "pt" or "pil" - num_images (`int`, *optional*, defaults to 1): - The number of images to generate Examples: Returns: diff --git a/predict.py b/predict.py index 921e8d4..8e5d321 100644 --- a/predict.py +++ b/predict.py @@ -97,7 +97,7 @@ def predict( description="Automatically adjust the output image size to be same as input image size. For editing and controlnet task, it can make sure the output image has the same size as input image leading to better performance", default=False, ), - num_images: int = Input(description="The number of images to generate for the given inputs", default=1, ge=1, le=10), + num_images_per_prompt: int = Input(description="The number of images to generate for the given inputs", default=1, ge=1, le=10), ) -> List[Path]: """Run a single prediction on the model""" if seed is None: @@ -106,7 +106,7 @@ def predict( input_images = [str(img) for img in [img1, img2, img3] if img is not None] output = [] - for i in range(num_images): + for i in range(num_images_per_prompt): output = self.pipe( prompt=prompt, input_images=None if len(input_images) == 0 else input_images, @@ -121,8 +121,7 @@ def predict( offload_model=offload_model, use_input_image_size_as_output=use_input_image_size_as_output, seed=seed, - max_input_image_size=max_input_image_size, - num_images=num_images, + max_input_image_size=max_input_image_size ) img = output[0] out_path = f"/tmp/out_{i}.png" From e9d5f163cb58479038032fbe0e50033518f359f6 Mon Sep 17 00:00:00 2001 From: asutermo Date: Fri, 10 Jan 2025 14:00:57 -0700 Subject: [PATCH 6/6] fix dumb mistake --- predict.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/predict.py b/predict.py index 8e5d321..78e9c9f 100644 --- a/predict.py +++ b/predict.py @@ -105,7 +105,7 @@ def predict( print(f"Using seed: {seed}") input_images = [str(img) for img in [img1, img2, img3] if img is not None] - output = [] + outputs = [] for i in range(num_images_per_prompt): output = self.pipe( prompt=prompt, @@ -126,5 +126,5 @@ def predict( img = output[0] out_path = f"/tmp/out_{i}.png" img.save(out_path) - output.append(Path(out_path)) - return output + outputs.append(Path(out_path)) + return outputs