Update README.md

README.md

@@ -67,33 +67,34 @@ instantir_path = f'./models/aggregator.pt'
 # load SDXL
 sdxl = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=torch.float16)
 
+# InstantIR pipeline
+pipe = InstantIRPipeline(
+    sdxl.vae, sdxl.text_encoder, sdxl.text_encoder_2, sdxl.tokenizer, sdxl.tokenizer_2,
+    sdxl.unet, sdxl.scheduler, feature_extractor=image_processor, image_encoder=image_encoder,
+)
+
 # load adapter
 image_proj_model = Resampler(
     embedding_dim=image_encoder.config.hidden_size,
     output_dim=sdxl.unet.config.cross_attention_dim,
 )
 init_adapter_in_unet(
-    sdxl.unet,
+    pipe.unet,
     image_proj_model,
     dcp_adapter,
 )
 
-pipe = InstantIRPipeline(
-    sdxl.vae, sdxl.text_encoder, sdxl.text_encoder_2, sdxl.tokenizer, sdxl.tokenizer_2,
-    sdxl.unet, sdxl.scheduler, feature_extractor=image_processor, image_encoder=image_encoder,
-)
-pipe.cuda()
-
 # load previewer lora
 pipe.prepare_previewers(previewer_lora_path)
-pipe.unet.to(dtype=torch.float16)
 pipe.scheduler = DDPMScheduler.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', subfolder="scheduler")
 lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)
+pipe.unet.to(dtype=torch.float16)
+pipe.to('cuda')
 
 # load aggregator weights
 pretrained_state_dict = torch.load(instantir_path)
 pipe.aggregator.load_state_dict(pretrained_state_dict)
-pipe.aggregator.to(dtype=torch.float16)
+pipe.aggregator.to(dtype=torch.float16, device=pipe.unet.device)
 ```
 
 Then, you can restore your broken images with: