Randomize output resolution - problem with writing custom script

[borderlineinteractive]

Hi,

I would like to implement a feature to add more variety to generations by randomizing the output image resolution. There are a few ways that I would like to implement, but I am already stuck with the simplest, first approach.

I would like to implement this as a custom script and I got this working to some extend with the following code:

import modules.scripts as scripts
import gradio as gr
import os
import random

from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state


class Script(scripts.Script):  

    def title(self):
        return "Randomize Resolution"

    def ui(self, is_img2img):
        is_enabled = gr.Checkbox(False, label="Randomize resolution")
        return [is_enabled]
      
    def show(self, is_img2img):
        
        return scripts.AlwaysVisible
        
    def before_process_batch(self, p, is_enabled, *args,**kwargs):
        if is_enabled:
            random.seed(p.seed+kwargs.get('batch_number'))
            choice=random.randint(1,7)
            print("choice:"+str(choice))
            match choice:
                case 1:
                    p.height=1024
                    p.width=1024
                case 2:
                    p.height=1152
                    p.width=896
                case 3:
                    p.height=896
                    p.width=1152
                case 4:
                    p.height=1216
                    p.width=832
                case 5:
                    p.height=832
                    p.width=1216
                case 6:
                    p.height=1344
                    p.width=768
                case 7:
                    p.height=768
                    p.width=1344
        print("height:"+str(p.height))
        print("width:"+str(p.width))

In principle the script works and it generates images with different resolutions for each image in a batch, and the randomization is linked to the seed and therefore reproducible. The problem with this code is however, that the height and width that I am setting in the "before_process_batch" function is not used in the current image, but instead in the following image. The result of this is that the PNG info of all generated images does not fit: It does not contain the resolution of the image itself, but instead the resolution of the next image that was processed afterwards.

Does anyone recognize the problem with this code?

Any help would be greatly appreciated.

[w-e-w]

stable-diffusion-webui/modules/processing.py

Lines 880 to 883 in e2b1990

	p.rng = rng.ImageRNG((opt_C, p.height // opt_f, p.width // opt_f), p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w)
	if p.scripts is not None:
	p.scripts.before_process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)

rng.ImageRNG is the class that dose the RNG noise, it is also the thing that defines the noize shape, the shape is configured on class __init__
as it is created before before_process_batch you can change the shape using p.height p.width inside before_process_batch

one method you could try setting the shape for the first itre in scripts.process(p)
and changing the shape for the next iteration at a later callback

a note of warning
as far as I'm aware most if not all internal processes and extensions does not expect the dimension to be changed from iteration to iteration, depending on the case you might call some compatibility issues

[borderlineinteractive]

So, if I understand you correctly, I could also just repeat the p.rng = rng.ImageRNG(... function again inside the custom script, and thereby the shape of the RNG noise should correspond correctly to the updated width and height. It is a bit of a hack, but I guess that the compute for the RNG is negligible in contrast to the overall compute. I will try if this works.

[w-e-w]

oh yea,I think you can do that
I didn't realize we store ImageRNG in p

I think you could just

from modules import scripts, processing

class Script(scripts.Script):
    def before_process_batch(*args, **kwargs)
       p.rng.shape = tuple(map(int, (processing.opt_C, new_height // processing.opt_f, new_width // opt_f)))

I guess that the compute for the RNG is negligible in contrast to the overall compute. I will try if this works.

my warning still stands, it's still possible that something might break if you change your resolution mid job
you should still set p.height .p.width but
some other extensions or maybe parts of the webui might look at these values and perform different actions based on them if I think the image is this size and didn't check the actual tensor then things might break
other things that are based on these values can also already be calculated and if so changing them in the job what else break stuff

for instance I'm suspecting that this will not go well with controlnet

that might also be some operations that requires all the image to be the same size

for example (this is only example I'm not sure if this is actually coded)
when stitching multiple image into a grid we might make the assumption that all images are the same size
and if you're not anymore things can break

---

if you can get yours to work then great
but to be honest if you wish to generate images at different size the safest way is to just do it in multiple jobs
if you wish to automate the process then the easiest way is to use API

[borderlineinteractive]

Thanks for the input. The script works now very nicely. I changed it a bit to automatically set the resolution based on the model (SDXL or SD15):

import modules.scripts as scripts
import gradio as gr
import os
import random

from modules import images, rng
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state



class Script(scripts.Script):  

# The title of the script. This is what will be displayed in the dropdown menu.
    def title(self):

        return "Randomize Resolution"

    def ui(self, is_img2img):
        is_enabled = gr.Checkbox(False, label="Randomize resolution")
        return [is_enabled]
      
    def show(self, is_img2img):
        
        return scripts.AlwaysVisible
        
    def before_process_batch(self, p, is_enabled, *args,**kwargs):
        
        if is_enabled:
            sd15_res_list = [(768,768),(768,512),(512,768),(768,576),(576,768),(912,512),(512,912)]
            sdxl_res_list = [(1024,1024),(1152,896),(896,1152),(1216,832),(832,1216),(1344,768),(768,1344)]
            opt_C = 4
            opt_f = 8
            random.seed(p.seed+kwargs.get('batch_number'))
            if p.sd_model.is_sdxl:
                res_tuple=random.choice(sdxl_res_list)
            else:
                res_tuple=random.choice(sd15_res_list)
            p.width=res_tuple[0]
            p.height=res_tuple[1]
            p.rng = rng.ImageRNG((opt_C, p.height // opt_f, p.width // opt_f), p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w)
            print("height:"+str(p.height))
            print("width:"+str(p.width))

[w-e-w]

nice

btw if you like to share your extension, you can have it listed on the extension index by makeing a PR to https://github.com/AUTOMATIC1111/stable-diffusion-webui-extensions