Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from diffusers.pipelines import StableDiffusionPipeline, StableDiffusionXLPipeline # pylint: disable=unused-import | |
| from modules import shared, scripts, processing, sd_models, devices | |
| """ | |
| This is a simpler template for script for SD.Next that implements a custom pipeline | |
| Items that can be added: | |
| - Any pipeline already in diffusers | |
| List of pipelines that can be directly used: <https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines> | |
| - Any pipeline for which diffusers definiotion exists and can be copied | |
| List of pipelines with community definitions: <https://github.com/huggingface/diffusers/tree/main/examples/community> | |
| - Any custom pipeline that you create | |
| Author:: | |
| - Your details | |
| Credits: | |
| - Link to original implementation and author | |
| Contributions: | |
| - Submit a PR on SD.Next GitHub repo to be included in /scripts | |
| - Before submitting a PR, make sure to test your script thoroughly and that it passes code quality checks | |
| Lint rules are part of SD.Next CI/CD pipeline | |
| > pip install ruff pylint | |
| > ruff scripts/example.py | |
| > pylint scriptts/example.py | |
| """ | |
| ## Config | |
| # script title | |
| title = 'Example' | |
| # is script available in txt2img tab | |
| txt2img = False | |
| # is script available in img2img tab | |
| img2img = False | |
| # is pipeline ok to run in pure latent mode without implicit conversions | |
| # recommended so entire ecosystem can be used as-is, but requires that latent is in format that sdnext can understand | |
| # some pipelines may not support this, in which case set to false and pipeline will implicitly do things like vae encode/decode on its own | |
| latent = True | |
| # base pipeline class from which this pipeline is derived, most commonly 'StableDiffusionPipeline' or 'StableDiffusionXLPipeline' | |
| pipeline_base = 'StableDiffusionPipeline' | |
| # class definition for this pipeline | |
| # for built-in diffuser pipelines, simply import it from diffusers.pipelines above | |
| # for example only, its set to same as base pipeline | |
| # for community pipelines, copy class definition from community source code | |
| # in which case only class definition code and required imports needs to be copied, not the entire source code | |
| pipeline_class = StableDiffusionPipeline | |
| # pipeline args values are defined in ui method below, here we need to define their exact names | |
| # they also have to be in the exact order as they are defined in ui | |
| # note: variable names should be exactly as defined in pipeline_class.__call__ method | |
| # if pipeline requires a param and its not provided, it will result in runtime error | |
| # if you provide param that is not defined by pipeline, sdnext will strip it | |
| params = ['test1', 'test2', 'test3', 'test4'] | |
| ### Script definition | |
| class Script(scripts.Script): | |
| def title(self): | |
| return title | |
| def show(self, is_img2img): | |
| if shared.backend == shared.Backend.DIFFUSERS: | |
| return img2img if is_img2img else txt2img | |
| return False | |
| # Define UI for pipeline | |
| def ui(self, _is_img2img): | |
| ui_controls = [] | |
| with gr.Row(): | |
| ui_controls.append(gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5, label="Test1")) | |
| ui_controls.append(gr.Slider(minimum=0, maximum=10, step=1, value=5, label="Test2")) | |
| with gr.Row(): | |
| ui_controls.append(gr.Checkbox(label="Test3", value=True)) | |
| with gr.Row(): | |
| ui_controls.append(gr.Textbox(label="Test4", value="", placeholder="enter text here")) | |
| with gr.Row(): | |
| gr.HTML('<a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl>Stable Diffusion SDXL pipeline docs</a>"') | |
| return ui_controls | |
| # Run pipeline | |
| def run(self, p: processing.StableDiffusionProcessing, *args): # pylint: disable=arguments-differ | |
| # prepare pipeline | |
| c = shared.sd_model.__class__.__name__ if shared.sd_model is not None else '' | |
| if c != pipeline_base: | |
| shared.log.warning(f'{title}: pipeline={c} required={pipeline_base}') | |
| return None | |
| orig_pipeline = shared.sd_model # backup current pipeline definition | |
| shared.sd_model = pipeline_class( # create new pipeline using currently loaded model which is always in `shared.sd_model` | |
| # different pipelines may need different init params, so you may need to change this | |
| # to see init params, see pipeline_class.__init__ method | |
| # if init params are incorrect you will also see a runtime error with unrecognized or missing params | |
| # for example: | |
| # > TypeError: StableDiffusionPipeline.__init__() missing 2 required positional arguments: 'safety_checker' and 'feature_extractor' | |
| vae = shared.sd_model.vae, | |
| text_encoder=shared.sd_model.text_encoder, | |
| tokenizer=shared.sd_model.tokenizer, | |
| unet=shared.sd_model.unet, | |
| scheduler=shared.sd_model.scheduler, | |
| safety_checker=shared.sd_model.safety_checker, | |
| feature_extractor=shared.sd_model.feature_extractor, | |
| ) | |
| sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) # copy options from original pipeline | |
| sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc. | |
| sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device | |
| shared.sd_model.to(dtype=devices.dtype) | |
| # if pipeline also needs a specific type, you can set it here, but not commonly needed | |
| # shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) | |
| # prepare params | |
| # all pipeline params go into p.task_args and are automatically handled by sdnext from there | |
| for i in range(len(args)): | |
| p.task_args[params[i]] = args[i] | |
| # you can also re-use existing params from `p` object if pipeline wants them, but under a different name | |
| # for example, if pipeline expects 'image' param, but you want to use 'init_images' instead which is what img2img tab uses | |
| # p.task_args['image'] = p.init_images[0] | |
| if not latent: | |
| p.task_args['output_type'] = 'np' | |
| shared.log.debug(f'{c}: args={p.task_args}') | |
| # if you need to run any preprocessing, this is the place to do it | |
| # run processing | |
| processed: processing.Processed = processing.process_images(p) | |
| # if you need to run any postprocessing, this is the place to do it | |
| # you dont need to handle saving, metadata, etc - sdnext will do it for you | |
| # restore original pipeline | |
| shared.sd_model = orig_pipeline | |
| return processed | |