Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add recursive option to upsample_imagefolder #154

Merged
merged 3 commits into from
Jan 26, 2023

Conversation

philgzl
Copy link
Contributor

@philgzl philgzl commented Jan 25, 2023

This PR implements a recursive option to RealESRGANModel.upsample_imagefolder such that all images in subdirectories inside the input directories are upscaled. Currently, only images one level down the input directory are upscaled. This is relevant to this project because StableDiffusionWalkPipeline.walk creates subdirectories of images inside the output directory.

Example: the user just created a video using the following code:

pipe.walk(
    prompts=["red cat", "green cat", "red cat"]
    output_dir="dreams",
    name="cats",
    ...
)

Now the directory dreams/ looks like this:

dreams
└── cats
    ├── cats_000000
    │   ├── frame000000.png
    │   ├── frame000001.png
    │   ├── frame000002.png
    │   └── ...
    ├── cats_000001
    │   ├── frame000000.png
    │   ├── frame000001.png
    │   ├── frame000002.png
    │   └── ...
    ├── cats.mp4
    └── prompt_config.json

Now the user wishes to upscale all the images. Calling RealESRGANModel.upsample_imagefolder(in_dir="dreams/cats", out_dir="path/to/output_dir") currently does nothing as there are no images inside dreams/cats/ (the images are in subfolders). The user currently needs to call the method for each subfolder inside dreams/cats/. This PR allows the user to provide recursive=True to upscale all images in subfolders without having to call the method for each subfolder. The output directory then presents the same structure as the input directory:

path/to/output_dir
├── cats_000000
│   ├── frame000000out.png
│   ├── frame000001out.png
│   ├── frame000002out.png
│   └── ...
└── cats_000001
    ├── frame000000out.png
    ├── frame000001out.png
    ├── frame000002out.png
    └── ...

Other additions:

  • Added a force option to allow overwriting of output files. Currently, output files are overwritten. With this PR, the default behavior is to skip if the output file already exists, unless force=True is provided.
  • Added logging. Logs tell if the image is skipped or upscaled, and the current progress (current image number over total number of images)

@philgzl philgzl changed the title Add recursive and option to upsampling_imagefolder Add recursive option to upsampling_imagefolder Jan 25, 2023
@philgzl philgzl changed the title Add recursive option to upsampling_imagefolder Add recursive option to upsample_imagefolder Jan 25, 2023
Copy link
Owner

@nateraw nateraw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one looks good to go, but will have to try myself before I can approve! will shoot over approving review + merge once I do.

Thanks for the contribution!!

Copy link
Owner

@nateraw nateraw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

image_paths = [x for x in generator if x.suffix.lower() in [".png", ".jpg", ".jpeg"]]
n_img = len(image_paths)
for i, image in enumerate(image_paths):
out_filepath = out_dir / (str(image.relative_to(in_dir).with_suffix('')) + suffix + outfile_ext)
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Little hard to read, but dissected it locally and I see why you had to do this. Maybe we think about writing this cleaner later (or not 😅 )

stable_diffusion_videos/upsampling.py Outdated Show resolved Hide resolved
@nateraw nateraw merged commit 79cb966 into nateraw:main Jan 26, 2023
@philgzl philgzl deleted the recursive-upsampling branch January 26, 2023 13:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants