ComfyUI Extension: Restart Sampling

Restart Sampling Introduction

ComfyUI_restart_sampling is an extension for the framework, designed to enhance the generative processes of AI models by implementing a novel sampling algorithm known as "Restart Sampling." This algorithm is based on the research paper "Restart Sampling for Improving Generative Processes" ().

The primary goal of this extension is to improve the balance between speed and quality in generative models, particularly those that solve differential equations, such as diffusion models. By using Restart Sampling, AI artists can achieve higher sample quality and faster sampling times, making it a valuable tool for creating high-quality AI-generated art efficiently.

How Restart Sampling Works

At its core, ComfyUI_restart_sampling leverages the Restart Sampling algorithm to optimize the generative process. Traditional samplers, like ODE-based and SDE-based samplers, have their own strengths and weaknesses. ODE-based samplers are fast but may plateau in performance, while SDE-based samplers produce higher quality samples but are slower.

Restart Sampling addresses these issues by introducing a mechanism that balances discretization errors and contraction. It periodically restarts the sampling process at specific intervals, which helps to reduce accumulated errors and improve the overall quality of the generated samples. Think of it as taking a break and resetting your progress to ensure you stay on the right track.

Restart Sampling Features

ComfyUI_restart_sampling offers several features to customize and optimize the sampling process:

KSampler With Restarts

This node includes all the inputs of a standard KSampler but adds a string widget for configuring Restart segments and a widget for the scheduler. It supports ODE samplers and is not intended for use with SDE samplers. The restart_scheduler is used for the denoising process during restart segments.

KSampler With Restarts (Simple)

This variant uses the same scheduler for both the KSampler and the restart segments, simplifying the configuration process.

KSampler With Restarts (Advanced)

This node includes all the inputs for an Advanced KSampler along with restart sampling inputs. It allows for more complex configurations but requires careful setup to avoid invalid segments.

KSampler With Restarts (Custom)

Similar to the Advanced version, but it takes a SAMPLER input like the built-in SamplerCustom node. This allows for more flexibility but requires compatibility checks with the Restart sampling process.

RestartScheduler

This node outputs sigmas with restart segments inserted and must be used with RestartSampler. It allows for custom sampling schedules and can take parameters for restart segments and schedules.

RestartSampler

Used in conjunction with RestartScheduler, this node arranges for the restart noise to be injected at the appropriate points and delegates the actual sampling to the supplied sampler.

Segments

Segments define the intervals at which the sampling process is restarted. They are specified as a sequence of arrays, each containing parameters like the number of restarts, steps, and time intervals. These segments can be customized to fine-tune the sampling process.

Chunked Mode

When enabled, the sampler is called with as many steps as possible up to the next segment, improving accuracy for certain samplers. This mode is particularly useful for samplers that rely on state from previous steps.

Troubleshooting Restart Sampling

Here are some common issues and their solutions:

Issue: Invalid Segments

Solution: Ensure that the segments are correctly defined and that the t_min and t_max values are appropriate. Invalid segments occur when the noise schedule does not align with the segment definitions.

Issue: Incompatible Samplers

Solution: Use only ODE-based samplers with Restart Sampling. SDE and UniPC samplers are not compatible and may cause errors.

Issue: No Output or Poor Quality Output

Solution: Check the configuration of the restart segments and the scheduler. Ensure that the environment variable COMFYUI_VERBOSE_RESTART_SAMPLING is set to 1 to get detailed information about the steps being run.

Learn More about Restart Sampling

For more information, you can refer to the following resources:

• These resources provide additional documentation, tutorials, and community support to help you get the most out of ComfyUI_restart_sampling.