q-future/Co-Instruct

1import os, yaml2import gradio as gr3import requests4import argparse5 6from PIL import Image7 8import numpy as np9import torch10from transformers import AutoModelForCausalLM11 12from huggingface_hub import hf_hub_download13 14 15## InstructIR Plugin ##16from insir_models import instructir17from insir_text.models import LanguageModel, LMHead18 19hf_hub_download(repo_id="marcosv/InstructIR", filename="im_instructir-7d.pt", local_dir="./")20hf_hub_download(repo_id="marcosv/InstructIR", filename="lm_instructir-7d.pt", local_dir="./")21 22CONFIG     = "eval5d.yml"23LM_MODEL   = "lm_instructir-7d.pt"24MODEL_NAME = "im_instructir-7d.pt"25 26def dict2namespace(config):27    namespace = argparse.Namespace()28    for key, value in config.items():29        if isinstance(value, dict):30            new_value = dict2namespace(value)31        else:32            new_value = value33        setattr(namespace, key, new_value)34    return namespace35 36 37# parse config file38with open(os.path.join(CONFIG), "r") as f:39    config = yaml.safe_load(f)40 41cfg = dict2namespace(config)42 43device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")44ir_model = instructir.create_model(input_channels =cfg.model.in_ch, width=cfg.model.width, enc_blks = cfg.model.enc_blks, 45                            middle_blk_num = cfg.model.middle_blk_num, dec_blks = cfg.model.dec_blks, txtdim=cfg.model.textdim)46ir_model = ir_model.to(device)47print ("IMAGE MODEL CKPT:", MODEL_NAME)48ir_model.load_state_dict(torch.load(MODEL_NAME, map_location="cpu"), strict=True)49 50os.environ["TOKENIZERS_PARALLELISM"] = "false"51LMODEL = cfg.llm.model52language_model = LanguageModel(model=LMODEL)53lm_head = LMHead(embedding_dim=cfg.llm.model_dim, hidden_dim=cfg.llm.embd_dim, num_classes=cfg.llm.nclasses)54lm_head = lm_head.to(device)55 56print("LMHEAD MODEL CKPT:", LM_MODEL)57lm_head.load_state_dict(torch.load(LM_MODEL, map_location="cpu"), strict=True)58 59def process_img(image, prompt=None):60    if prompt is None:61        prompt = chat("How to improve the quality of the image?", [], image, None, None, None)62        prompt += "Please help me improve its quality!"63        print(prompt)64    img = np.array(image)65    img = img / 255.66    img = img.astype(np.float32)67    y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)68 69    lm_embd = language_model(prompt)70    lm_embd = lm_embd.to(device)71 72    with torch.no_grad():73        text_embd, deg_pred = lm_head(lm_embd)74        x_hat = ir_model(y, text_embd)75 76    restored_img = x_hat.squeeze().permute(1,2,0).clamp_(0, 1).cpu().detach().numpy()77    restored_img = np.clip(restored_img, 0. , 1.)78 79    restored_img = (restored_img * 255.0).round().astype(np.uint8)  # float32 to uint880    return Image.fromarray(restored_img) #(image, Image.fromarray(restored_img))81 82## InstructIR Plugin ##83model = AutoModelForCausalLM.from_pretrained("q-future/co-instruct-preview", 84                                             trust_remote_code=True, 85                                             torch_dtype=torch.float16, 86                                             attn_implementation="eager",87                                             device_map={"":"cuda:0"})88 89def chat(message, history, image_1, image_2, image_3, image_4):90    print(history)91    if history:92        if image_1 is not None and image_2 is None:93            past_message = "USER: The input image: <|image|>" + history[0][0] + " ASSISTANT:" + history[0][1]94            for i in range((len(history) - 1)):95                past_message += "USER:" +history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"96            message = past_message + "USER:" + message + " ASSISTANT:"97            images = [image_1]98        if image_1 is not None and image_2 is not None:99            if image_3 is None:100                past_message = "USER: The first image: <|image|>\nThe second image: <|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"101                for i in range((len(history) - 1)):102                    past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"103                message = past_message + "USER:" + message + " ASSISTANT:"104                images = [image_1, image_2]105            else:106                if image_4 is None:107                    past_message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"108                    for i in range((len(history) - 1)):109                        past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"110                    message = past_message + "USER:" + message + " ASSISTANT:"111                    images = [image_1, image_2, image_3]112                else:113                    past_message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>\nThe fourth image:<|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"114                    for i in range((len(history) - 1)):115                        past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"116                    message = past_message + "USER:" + message + " ASSISTANT:"117                    images = [image_1, image_2, image_3, image_4]118    else:  119        if image_1 is not None and image_2 is None:120            message = "USER: The input image: <|image|>" + message + " ASSISTANT:"121            images = [image_1]122        if image_1 is not None and image_2 is not None:123            if image_3 is None:124                message = "USER: The first image: <|image|>\nThe second image: <|image|>" + message + " ASSISTANT:"125                images = [image_1, image_2]126            else:127                if image_4 is None:128                    message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>" + message + " ASSISTANT:"129                    images = [image_1, image_2, image_3]130                else:131                    message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>\nThe fourth image:<|image|>" + message + " ASSISTANT:"132                    images = [image_1, image_2, image_3, image_4]133    134    print(message)135    136    return model.tokenizer.batch_decode(model.chat(message, images, max_new_tokens=600).clamp(0, 100000))[0].split("ASSISTANT:")[-1]137 138#### Image,Prompts examples139examples = [140    ["Which part of the image is relatively clearer, the upper part or the lower part? Please analyze in details.", "examples/sausage.jpg", None],141    ["Which image is noisy, and which one is with motion blur? Please analyze in details.", "examples/211.jpg", "examples/frog.png"],142            ["What is the problem in this image, and how to fix it? Please answer my questions one by one.", "examples/lol_748.png", None],143]144 145#<h1 align="center"><a href="https://github.com/Q-Future/Q-Instruct"><img src="https://github.com/Q-Future/Q-Instruct/blob/main/q_instruct_logo.png?raw=true", alt="Q-Instruct (mPLUG-Owl-2)" border="0" style="margin: 0 auto; height: 85px;" /></a> </h1>146 147 148title = "Co-Instruct-Plus🧑‍🏫🖌️"149with gr.Blocks(title="Co-Instruct-Plus🧑‍🏫🖌️") as demo:150    title_markdown = ("""151 152<h1 align="center"><a href="https://github.com/Q-Future/Co-Instruct"><img src="https://raw.githubusercontent.com/Q-Future/Co-Instruct/main/co-instruct.png", alt="Co-Instruct" border="0" style="margin: 0 auto; height: 85px;" /></a> </h1>153 154<div align="center">Built upon <strong>Q-Instruct: Improving Low-level Visual Abilities for Multi-modality Foundation Models (CVPR 2024)</strong></div>155 156<div align="center">Built upon the Upgraded Version, Co-Instruct, supporting up to 4 images: <strong>Towards Open-ended Visual Quality Comparison (Arxiv 2024)</strong></div>   157 158<div align="center">We also support <a href='https://huggingface.co/marcosv/InstructIR'>InstructIR</a> as PLUGIN to restore image!</div>159<h5 align="center"> Please find our more accurate visual scoring demo on <a href='https://huggingface.co/spaces/teowu/OneScorer'>[OneScorer]</a> (Q-Align)!</h2>160<div align="center">161    <div style="display:flex; gap: 0.25rem;" align="center">162    <strong>Q-Instruct Resources:</strong>163        <a href='https://github.com/Q-Future/Q-Instruct'><img src='https://img.shields.io/badge/Github-Code-blue'></a>164        <a href="https://Q-Instruct.github.io/Q-Instruct/fig/Q_Instruct_v0_1_preview.pdf"><img src="https://img.shields.io/badge/Technical-Report-red"></a>165        <a href='https://github.com/Q-Future/Q-Instruct/stargazers'><img src='https://img.shields.io/github/stars/Q-Future/Q-Instruct.svg?style=social'></a>166    </div>167</div>168<div align="center">169    <div style="display:flex; gap: 0.25rem;" align="center">170    <strong>Co-Instruct Resources:</strong>171        <a href='https://github.com/Q-Future/Co-Instruct'><img src='https://img.shields.io/badge/Github-Code-blue'></a>172        <a href="https://arxiv.org/pdf/2402.16641.pdf"><img src="https://img.shields.io/badge/Technical-Report-red"></a>173        <a href='https://github.com/Q-Future/Co-Instruct/stargazers'><img src='https://img.shields.io/github/stars/Q-Future/Co-Instruct.svg?style=social'></a>174    </div>175</div>176""")177    gr.Markdown(title_markdown)178    with gr.Row():179            input_img_1 = gr.Image(type='pil', label="Image 1 (First image)")180            input_img_2 = gr.Image(type='pil', label="Image 2 (Second image)")181            input_img_3 = gr.Image(type='pil', label="Image 3 (Third image)")182            input_img_4 = gr.Image(type='pil', label="Image 4 (Fourth image)")183    with gr.Row():184        with gr.Column(scale=2):185            gr.ChatInterface(fn = chat, additional_inputs=[input_img_1, input_img_2, input_img_3, input_img_4], theme="Soft", examples=examples)186        with gr.Column(scale=1):187            input_image_ir = gr.Image(type="pil", label="Image for Auto Restoration")188            output_image_ir = gr.Image(type="pil", label="Output of Auto Restoration")189            gr.Interface(190                fn=process_img,191                inputs=[input_image_ir],192                outputs=[output_image_ir],193                examples=["examples/gopro.png", "examples/noise50.png", "examples/lol_748.png"],194            )195    demo.launch(share=True)