nateraw / stable-diffusion-videos

Create 🔥 videos with Stable Diffusion by exploring the latent space and morphing between text prompts
Apache License 2.0
4.34k stars 413 forks source link

Add `recursive` option to `upsample_imagefolder` #154

Closed philgzl closed 1 year ago

philgzl commented 1 year ago

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: