Update README.md

README.md

@@ -9,8 +9,8 @@ pipeline_tag: image-to-image
 # InstantIR Model Card
 <div style="display: flex; gap: 10px; align-items: center; justify-content: center; height: auto;">
 <a href='https://arxiv.org/abs/2410.06551'><img src='https://img.shields.io/badge/paper-arXiv-b31b1b.svg' style="height: 24px;"></a>
-<a href='https://jy-joy.github.io/InstantIR'><img src='https://img.shields.io/badge/project-Website-informational' style="height: 24px;"></a>
-<a href='https://github.com/JY-Joy/InstantIR'><img src='https://img.shields.io/badge/code-Github-gray' style="height: 24px;"></a>
+<a href='https://jy-joy.github.io/InstantIR'><img src='https://img.shields.io/badge/project-Website-green' style="height: 24px;"></a>
+<a href='https://github.com/JY-Joy/InstantIR'><img src='https://img.shields.io/badge/code-Github-informational' style="height: 24px;"></a>
 </div>
 
 > **InstantIR** is a novel single-image restoration model designed to resurrect your damaged images, delivering extrem-quality yet realistic details. You can further boost **InstantIR** performance with additional text prompts, even achieve customized editing!
@@ -46,55 +46,41 @@ hf_hub_download(repo_id="InstantX/InstantIR", filename="models/previewer_lora_we
 import torch
 from PIL import Image
 
-from diffusers import DDPMScheduler, StableDiffusionXLPipeline
+from diffusers import DDPMScheduler
 from schedulers.lcm_single_step_scheduler import LCMSingleStepScheduler
 
 from transformers import AutoImageProcessor, AutoModel
 
-from module.ip_adapter.utils import init_adapter_in_unet
-from module.ip_adapter.resampler import Resampler
+from module.ip_adapter.utils import load_adapter_to_pipe
 from pipelines.sdxl_instantir import InstantIRPipeline
 
-# prepare 'dinov2'
-image_encoder = AutoModel.from_pretrained('facebook/dinov2-large')
-image_processor = AutoImageProcessor.from_pretrained('facebook/dinov2-large')
-
 # prepare models under ./checkpoints
 dcp_adapter = f'./models/adapter.pt'
 previewer_lora_path = f'./models'
 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 pretrained models
+pipe = InstantIRPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=torch.float16)
 
 # 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(
-    pipe.unet,
-    image_proj_model,
+load_adapter_to_pipe(
+    pipe,
     dcp_adapter,
+    image_encoder_or_path = 'facebook/dinov2-large',
 )
 
 # load previewer lora
 pipe.prepare_previewers(previewer_lora_path)
 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, device=pipe.unet.device)
+
+# send to GPU and fp16
+pipe.to(dtype=torch.float16)
+pipe.to('cuda')
 ```
 
 Then, you can restore your broken images with: