only train and infer stage 2 model

[tzayuan]

Hi, @0x3f3f3f3fun I want to train & inference only stage 2 model(stage 1 will implement alone by another model and generate restorationed image), if you could provide me a command structures?(how to only fine-tune stage 2 model and save it, and how to inference only stage 2 model) Thanks!

BRs, tzayuan

[0x3f3f3f3fun]

Hi :) For inference only stage 2 model, you can take the following steps:

1. Implement a custom pipeline without any restoration module in utils/helpers.py.

   class CustomPipeline(Pipeline):
   
   def __init__(self, cldm: ControlLDM, diffusion: Diffusion, cond_fn: Optional[Guidance], device: str) -> None:
       super().__init__(None, cldm, diffusion, cond_fn, device)
   
   @count_vram_usage
   def run_stage1(self, lq: torch.Tensor) -> torch.Tensor:
       # In our experiments, the output of restoration module (a.k.a the condition for subsequent IRControlNet)
       # will be resized to a resolution >= 512, since both pretrained SD2 and IRControlNet are trained on 512x512.
       # Here we directly use the lq as condition.
       if min(lq.shape[2:]) < 512:
           clean = resize_short_edge_to(lq, size=512)
       return clean

2. Implement a custom inference loop in utils/inference.py.

   class CustomInferenceLoop(InferenceLoop):
   
   @count_vram_usage
   def init_stage1_model(self) -> None:
       # Nothing to do.
       pass
   
   def init_pipeline(self) -> None:
       # Instantiate our custom pipeline.
       self.pipeline = CustomPipeline(self.cldm, self.diffusion, self.cond_fn, self.args.device)

3. Add the custom pipeline to inference.py as an option.

   def main():
   args = parse_args()
   args.device = check_device(args.device)
   set_seed(args.seed)
   if args.version == "v1":
       V1InferenceLoop(args).run()
   else:
       supported_tasks = {
           "sr": BSRInferenceLoop,
           "dn": BIDInferenceLoop,
           "fr": BFRInferenceLoop,
           "fr_bg": UnAlignedBFRInferenceLoop,
           "custom": CustomInferenceLoop
       }
       supported_tasks[args.task](args).run()
       print("done!")

   The provided code is directly wrote in the GitHub comment console and has not been tested, but I think it will work.

As for finetuning stage2 model, there are two ways to do it:

• Apply the restoration model to low-quality images off-line, and write a new dataset class to load the gt-condition pairs.
• Copy the model denifition to model directory, load it in train_stage2.py and replace SwinIR with your restoration model.

The second way is more recommended.