Upload image_size_randomize_v2.py
Browse files- image_size_randomize_v2.py +117 -0
    	
        image_size_randomize_v2.py
    ADDED
    
    | @@ -0,0 +1,117 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            import random
         | 
| 2 | 
            +
            import modules.scripts as scripts
         | 
| 3 | 
            +
            import gradio as gr
         | 
| 4 | 
            +
            import os
         | 
| 5 | 
            +
             | 
| 6 | 
            +
            from modules import images
         | 
| 7 | 
            +
            from modules.processing import process_images, Processed
         | 
| 8 | 
            +
            from modules.shared import opts, cmd_opts, state
         | 
| 9 | 
            +
             | 
| 10 | 
            +
            class RandomImageSizeScript(scripts.Script):
         | 
| 11 | 
            +
                def __init__(self):
         | 
| 12 | 
            +
                    super().__init__()
         | 
| 13 | 
            +
                    self.index = 0  # Round-Robin mode index
         | 
| 14 | 
            +
                    self.swap_round = False  # Controls width-height swap per round
         | 
| 15 | 
            +
             | 
| 16 | 
            +
                def title(self):
         | 
| 17 | 
            +
                    return "Image Size Randomize"
         | 
| 18 | 
            +
             | 
| 19 | 
            +
                def ui(self, is_txt2img):
         | 
| 20 | 
            +
                    with gr.Accordion("Image Size Randomize Settings", open=True):
         | 
| 21 | 
            +
                        with gr.Group():
         | 
| 22 | 
            +
                            gr.Markdown("### Size")
         | 
| 23 | 
            +
                            gr.HTML("<p>Enter image sizes in the format 'width,height', separated by new lines.</p>")
         | 
| 24 | 
            +
                            text_sizes = gr.Textbox(
         | 
| 25 | 
            +
                                label='Image Sizes',
         | 
| 26 | 
            +
                                placeholder='Enter sizes, e.g.:\n1024,512\n1024,768\n1024,1024',
         | 
| 27 | 
            +
                                elem_id=self.elem_id("text_sizes"),
         | 
| 28 | 
            +
                                lines=10
         | 
| 29 | 
            +
                            )
         | 
| 30 | 
            +
             | 
| 31 | 
            +
                        with gr.Group():
         | 
| 32 | 
            +
                            gr.Markdown("### Mode")
         | 
| 33 | 
            +
                            mode = gr.Radio(
         | 
| 34 | 
            +
                                choices=["Random", "Round-Robin"],
         | 
| 35 | 
            +
                                value="Random",
         | 
| 36 | 
            +
                                label="Selection Mode",
         | 
| 37 | 
            +
                                elem_id=self.elem_id("mode"),
         | 
| 38 | 
            +
                                info="Choose whether to select sizes randomly or in a round-robin fashion."
         | 
| 39 | 
            +
                            )
         | 
| 40 | 
            +
             | 
| 41 | 
            +
                            random_swap_text = gr.Checkbox(
         | 
| 42 | 
            +
                                label='Randomly swap width and height',
         | 
| 43 | 
            +
                                elem_id=self.elem_id("random_swap_text"),
         | 
| 44 | 
            +
                                visible=True
         | 
| 45 | 
            +
                            )
         | 
| 46 | 
            +
                            
         | 
| 47 | 
            +
                            sequential_swap = gr.Checkbox(
         | 
| 48 | 
            +
                                label='Swap width and height every round',
         | 
| 49 | 
            +
                                elem_id=self.elem_id("sequential_swap"),
         | 
| 50 | 
            +
                                visible=False
         | 
| 51 | 
            +
                            )
         | 
| 52 | 
            +
             | 
| 53 | 
            +
                            def update_visibility(selected_mode):
         | 
| 54 | 
            +
                                if selected_mode == "Random":
         | 
| 55 | 
            +
                                    return {
         | 
| 56 | 
            +
                                        random_swap_text: gr.update(visible=True),
         | 
| 57 | 
            +
                                        sequential_swap: gr.update(visible=False)
         | 
| 58 | 
            +
                                    }
         | 
| 59 | 
            +
                                else:
         | 
| 60 | 
            +
                                    return {
         | 
| 61 | 
            +
                                        random_swap_text: gr.update(visible=False),
         | 
| 62 | 
            +
                                        sequential_swap: gr.update(visible=True)
         | 
| 63 | 
            +
                                    }
         | 
| 64 | 
            +
             | 
| 65 | 
            +
                            mode.change(
         | 
| 66 | 
            +
                                fn=update_visibility,
         | 
| 67 | 
            +
                                inputs=mode,
         | 
| 68 | 
            +
                                outputs=[random_swap_text, sequential_swap]
         | 
| 69 | 
            +
                            )
         | 
| 70 | 
            +
             | 
| 71 | 
            +
                    return [text_sizes, random_swap_text, mode, sequential_swap]
         | 
| 72 | 
            +
             | 
| 73 | 
            +
                def run(self, p, text_sizes, random_swap_text, mode, sequential_swap):
         | 
| 74 | 
            +
                    sizes = []
         | 
| 75 | 
            +
             | 
| 76 | 
            +
                    # Parse and validate input sizes
         | 
| 77 | 
            +
                    if text_sizes:
         | 
| 78 | 
            +
                        try:
         | 
| 79 | 
            +
                            sizes = [tuple(map(int, line.strip().split(','))) for line in text_sizes.strip().split('\n') if line.strip()]
         | 
| 80 | 
            +
                        except ValueError:
         | 
| 81 | 
            +
                            print(f"Invalid size format detected. Ensure all sizes are in the format 'width,height'.")
         | 
| 82 | 
            +
                    
         | 
| 83 | 
            +
                    if sizes:
         | 
| 84 | 
            +
                        try:
         | 
| 85 | 
            +
                            if mode == "Random":
         | 
| 86 | 
            +
                                width, height = random.choice(sizes)
         | 
| 87 | 
            +
                                if random_swap_text and random.choice([True, False]):
         | 
| 88 | 
            +
                                    width, height = height, width
         | 
| 89 | 
            +
                            else:  # Round-Robin mode
         | 
| 90 | 
            +
                                width, height = sizes[self.index]
         | 
| 91 | 
            +
             | 
| 92 | 
            +
                                # Handle dynamic size changes and index resets
         | 
| 93 | 
            +
                                if self.index >= len(sizes):
         | 
| 94 | 
            +
                                    self.index = 0  # Reset index if out of bounds
         | 
| 95 | 
            +
                                    width, height = sizes[self.index]
         | 
| 96 | 
            +
             | 
| 97 | 
            +
                                if self.index == 0 and sequential_swap:
         | 
| 98 | 
            +
                                    self.swap_round = not self.swap_round
         | 
| 99 | 
            +
             | 
| 100 | 
            +
                                if sequential_swap and self.swap_round:
         | 
| 101 | 
            +
                                    width, height = height, width
         | 
| 102 | 
            +
             | 
| 103 | 
            +
                                # Update index for the next round
         | 
| 104 | 
            +
                                self.index = (self.index + 1) % len(sizes)
         | 
| 105 | 
            +
             | 
| 106 | 
            +
                            # Set dimensions for image processing
         | 
| 107 | 
            +
                            p.width, p.height = width, height
         | 
| 108 | 
            +
             | 
| 109 | 
            +
                        except IndexError as e:
         | 
| 110 | 
            +
                            print(f"Error: {e}. The size list may have been modified unexpectedly.")
         | 
| 111 | 
            +
                            self.index = 0  # Reset index and recover from error
         | 
| 112 | 
            +
             | 
| 113 | 
            +
                    else:
         | 
| 114 | 
            +
                        print("No valid sizes found. Please check the input format.")
         | 
| 115 | 
            +
             | 
| 116 | 
            +
                    proc = process_images(p)
         | 
| 117 | 
            +
                    return proc
         |