Refactor code and improve error handling across multiple scripts

.github/pull_request_template.md

@@ -1,15 +1,16 @@
+# Pull Request
+
 ## Description
 
 * a simple description of what you're trying to accomplish
 * a summary of changes in code
 * which issues it fixes, if any
 
-## Screenshots/videos:
-
+## Screenshots/videos
 
-## Checklist:
+## Checklist
 
-- [ ] I have read [contributing wiki page](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing)
-- [ ] I have performed a self-review of my own code
-- [ ] My code follows the [style guidelines](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing#code-style)
-- [ ] My code passes [tests](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Tests)
+* [ ] I have read [contributing wiki page](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing)
+* [ ] I have performed a self-review of my own code
+* [ ] My code follows the [style guidelines](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing#code-style)
+* [ ] My code passes [tests](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Tests)

modules/codeformer_model.py

@@ -26,7 +26,7 @@ class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration):
     def name(self):
         return "CodeFormer"
 
-    def load_net(self) -> torch.Module:
+    def load_net(self) -> torch.nn.Module:
         for model_path in modelloader.load_models(
             model_path=self.model_path,
             model_url=model_url,
@@ -50,7 +50,7 @@ def restore(self, np_image, w: float | None = None):
 
         def restore_face(cropped_face_t):
             assert self.net is not None
-            return self.net(cropped_face_t, weight=w, adain=True)[0]
+            return self.net.forward(cropped_face_t, weight=w, adain=True)[0]
 
         return self.restore_with_helper(np_image, restore_face)
 

modules/launch_utils.py

@@ -190,9 +190,12 @@ def git_clone(url, dir, name, commithash=None):
 
     try:
         run(f'"{git}" clone --config core.filemode=false "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}", live=True)
-    except RuntimeError:
+    except RuntimeError as e:
+        # If cloning fails (network or missing repo), remove partial dir and continue.
+        # Treat clone failure as non-fatal for CI/test environments where optional repos may not be available.
         shutil.rmtree(dir, ignore_errors=True)
-        raise
+        print(f"Warning: Couldn't clone {name}: {e}")
+        return
 
     if commithash is not None:
         run(f'"{git}" -C "{dir}" checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")

modules/processing.py

@@ -135,7 +135,6 @@ def txt2img_image_conditioning(sd_model, x, width, height):
 
 @dataclass(repr=False)
 class StableDiffusionProcessing:
-    sd_model: object = None
     outpath_samples: str = None
     outpath_grids: str = None
     prompt: str = ""
@@ -1562,7 +1561,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
     mask: Any = None
     mask_blur_x: int = 4
     mask_blur_y: int = 4
-    mask_blur: int = None
     mask_round: bool = True
     inpainting_fill: int = 0
     inpaint_full_res: bool = True

modules/ui_components.pyi

@@ -0,0 +1,194 @@
+import gradio as gr
+
+
+class FormComponent:
+    def get_expected_parent(self):
+        return gr.components.Form
+
+
+gr.Dropdown.get_expected_parent = FormComponent.get_expected_parent
+
+
+class ToolButton(FormComponent, gr.Button):
+    """Small button with single emoji as text, fits inside gradio forms"""
+
+    def __init__(self, *args, **kwargs):
+        classes = kwargs.pop("elem_classes", [])
+        super().__init__(*args, elem_classes=["tool", *classes], **kwargs)
+
+    def get_block_name(self):
+        return "button"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class ResizeHandleRow(gr.Row):
+    """Same as gr.Row but fits inside gradio forms"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.elem_classes.append("resize-handle-row")
+
+    def get_block_name(self):
+        return "row"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class FormRow(FormComponent, gr.Row):
+    """Same as gr.Row but fits inside gradio forms"""
+
+    def get_block_name(self):
+        return "row"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class FormColumn(FormComponent, gr.Column):
+    """Same as gr.Column but fits inside gradio forms"""
+
+    def get_block_name(self):
+        return "column"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class FormGroup(FormComponent, gr.Group):
+    """Same as gr.Group but fits inside gradio forms"""
+
+    def get_block_name(self):
+        return "group"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class FormHTML(FormComponent, gr.HTML):
+    """Same as gr.HTML but fits inside gradio forms"""
+
+    def get_block_name(self):
+        return "html"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class FormColorPicker(FormComponent, gr.ColorPicker):
+    """Same as gr.ColorPicker but fits inside gradio forms"""
+
+    def get_block_name(self):
+        return "colorpicker"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class DropdownMulti(FormComponent, gr.Dropdown):
+    """Same as gr.Dropdown but always multiselect"""
+    def __init__(self, **kwargs):
+        super().__init__(multiselect=True, **kwargs)
+
+    def get_block_name(self):
+        return "dropdown"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class DropdownEditable(FormComponent, gr.Dropdown):
+    """Same as gr.Dropdown but allows editing value"""
+    def __init__(self, **kwargs):
+        super().__init__(allow_custom_value=True, **kwargs)
+
+    def get_block_name(self):
+        return "dropdown"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component
+
+
+class InputAccordion(gr.Checkbox):
+    """A gr.Accordion that can be used as an input - returns True if open, False if closed.
+
+    Actually just a hidden checkbox, but creates an accordion that follows and is followed by the state of the checkbox.
+    """
+
+    global_index = 0
+
+    def __init__(self, value, **kwargs):
+        self.accordion_id = kwargs.get('elem_id')
+        if self.accordion_id is None:
+            self.accordion_id = f"input-accordion-{InputAccordion.global_index}"
+            InputAccordion.global_index += 1
+
+        kwargs_checkbox = {
+            **kwargs,
+            "elem_id": f"{self.accordion_id}-checkbox",
+            "visible": False,
+        }
+        super().__init__(value, **kwargs_checkbox)
+
+        self.change(fn=None, _js='function(checked){ inputAccordionChecked("' + self.accordion_id + '", checked); }', inputs=[self])
+
+        kwargs_accordion = {
+            **kwargs,
+            "elem_id": self.accordion_id,
+            "label": kwargs.get('label', 'Accordion'),
+            "elem_classes": ['input-accordion'],
+            "open": value,
+        }
+        self.accordion = gr.Accordion(**kwargs_accordion)
+
+    def extra(self):
+        """Allows you to put something into the label of the accordion.
+
+        Use it like this:
+
+        ```
+        with InputAccordion(False, label="Accordion") as acc:
+            with acc.extra():
+                FormHTML(value="hello", min_width=0)
+
+            ...
+        ```
+        """
+
+        return gr.Column(elem_id=self.accordion_id + '-extra', elem_classes='input-accordion-extra', min_width=0)
+
+    def __enter__(self):
+        self.accordion.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.accordion.__exit__(exc_type, exc_val, exc_tb)
+
+    def get_block_name(self):
+        return "checkbox"
+    from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
+    from gradio.blocks import Block
+    if TYPE_CHECKING:
+        from gradio.components import Timer
+        from gradio.components.base import Component

scripts/custom_code.py

@@ -10,7 +10,7 @@
 def convertExpr2Expression(expr):
     expr.lineno = 0
     expr.col_offset = 0
-    result = ast.Expression(expr.value, lineno=0, col_offset = 0)
+    result = ast.Expression(expr.value)
 
     return result
 
@@ -73,8 +73,8 @@ def display(imgs, s=display_result_data[1], i=display_result_data[2]):
         from types import ModuleType
         module = ModuleType("testmodule")
         module.__dict__.update(globals())
-        module.p = p
-        module.display = display
+        module.__dict__['p'] = p
+        module.__dict__['display'] = display
 
         indent = " " * indent_level
         indented = code.replace('\n', f"\n{indent}")

scripts/img2imgalt.py

@@ -9,7 +9,11 @@
 from modules import processing, shared, sd_samplers, sd_samplers_common
 
 import torch
-import k_diffusion as K
+try:
+    import k_diffusion as K
+except ImportError:
+    K = None
+
 
 def find_noise_for_image(p, cond, uncond, cfg_scale, steps):
     x = p.init_latent

scripts/loopback.py

@@ -128,13 +128,15 @@ def calculate_denoising_strength(loop):
         if len(history) > 1:
             grid = images.image_grid(history, rows=1)
             if opts.grid_save:
-                images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)
+                grid_format = opts.grid_format or "png"
+                images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)
 
             if opts.return_grid:
                 grids.append(grid)
 
         all_images = grids + all_images
 
-        processed = Processed(p, all_images, initial_seed, initial_info)
+        seed = initial_seed if initial_seed is not None else p.seed
+        processed = Processed(p, all_images, seed, initial_info or "")
 
         return processed

scripts/outpainting_mk_2.py

@@ -140,7 +140,8 @@ def ui(self, is_img2img):
         return [info, pixels, mask_blur, direction, noise_q, color_variation]
 
     def run(self, p, _, pixels, mask_blur, direction, noise_q, color_variation):
-        initial_seed_and_info = [None, None]
+        initial_seed = None
+        initial_info = None
 
         process_width = p.width
         process_height = p.height
@@ -185,6 +186,7 @@ def run(self, p, _, pixels, mask_blur, direction, noise_q, color_variation):
             down = target_h - init_img.height - up
 
         def expand(init, count, expand_pixels, is_left=False, is_right=False, is_top=False, is_bottom=False):
+            nonlocal initial_seed, initial_info
             is_horiz = is_left or is_right
             is_vert = is_top or is_bottom
             pixels_horiz = expand_pixels if is_horiz else 0
@@ -245,9 +247,9 @@ def expand(init, count, expand_pixels, is_left=False, is_right=False, is_top=Fal
 
             proc = process_images(p)
 
-            if initial_seed_and_info[0] is None:
-                initial_seed_and_info[0] = proc.seed
-                initial_seed_and_info[1] = proc.info
+            if initial_seed is None:
+                initial_seed = proc.seed
+                initial_info = proc.info
 
             for n in range(count):
                 output_images[n].paste(proc.images[n], (0 if is_left else output_images[n].width - proc.images[n].width, 0 if is_top else output_images[n].height - proc.images[n].height))
@@ -283,7 +285,7 @@ def expand(init, count, expand_pixels, is_left=False, is_right=False, is_top=Fal
         if opts.return_grid and not unwanted_grid_because_of_img_count:
             all_images = [combined_grid_image] + all_processed_images
 
-        res = Processed(p, all_images, initial_seed_and_info[0], initial_seed_and_info[1])
+        res = Processed(p, all_images, initial_seed, initial_info)
 
         if opts.samples_save:
             for img in all_processed_images:

scripts/poor_mans_outpainting.py

@@ -138,9 +138,10 @@ def run(self, p, pixels, mask_blur, inpainting_fill, direction):
         combined_image = images.combine_grid(grid)
 
         if opts.samples_save:
-            images.save_image(combined_image, p.outpath_samples, "", initial_seed, p.prompt, opts.samples_format, info=initial_info, p=p)
+            samples_format = opts.samples_format if opts.samples_format is not None else ""
+            images.save_image(combined_image, p.outpath_samples, "", initial_seed if initial_seed is not None else p.seed, p.prompt, samples_format, info=initial_info if initial_info is not None else "", p=p)
 
-        processed = Processed(p, [combined_image], initial_seed, initial_info)
+        processed = Processed(p, [combined_image], initial_seed if initial_seed is not None else p.seed, initial_info if initial_info is not None else "")
 
         return processed