mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
133 lines
6.0 KiB
Python
133 lines
6.0 KiB
Python
import gradio as gr
|
|
from diffusers.pipelines import StableDiffusionPipeline, StableDiffusionXLPipeline # pylint: disable=unused-import
|
|
from modules import shared, scripts, processing, sd_models
|
|
|
|
"""
|
|
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
|
|
|
|
# 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_pipepeline = 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,
|
|
)
|
|
if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
|
|
shared.sd_model.to(shared.device) # move pipeline if needed, but don't touch if its under automatic managment
|
|
sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc.
|
|
|
|
# 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_pipepeline
|
|
return processed
|