add moondream

CHANGELOG.md

@@ -5,7 +5,7 @@
 - EDM samplers for Playground require `diffusers==0.27.0`
 - StableCascade requires diffusers `kashif/diffusers.git@wuerstchen-v3`
 
-## Update for 2024-03-03
+## Update for 2024-03-05
 
 - [Playground v2.5](https://huggingface.co/playgroundai/playground-v2.5-1024px-aesthetic)
   - new model version from Playground: based on SDXL, but with some cool new concepts
@@ -15,17 +15,20 @@
   - another very fast & light sd-xl model where original unet was compressed and distilled to 54% of original size  
   - download using networks -> reference
   - *note* to download fp16 variant (recommended), set settings -> diffusers -> preferred model variant  
-- **Image2Video**
-  - new module for creating videos from images  
-  - simply enable from *img2img -> scripts -> image2video*  
-  - based on [VGen](https://huggingface.co/ali-vilab/i2vgen-xl)  
-- **VQA** visual question & answer in interrogate  
-  - with support for multiple variations of base models: *GIT, BLIP, ViLT, PIX*
+- **Visual Query** visual query & answer in process tab  
+  - ask your questions, e.g. "describe the image", "what is behind the subject", "what are predominant colors of the image?"
+  - primary model is [moondream2](https://github.com/vikhyat/moondream), a *tiny* 1.86B vision language model  
+    (its still 3.7GB in size, so not really tiny)  
+  - additional support for multiple variations of several base models: *GIT, BLIP, ViLT, PIX*
 - **Second Pass / Refine**
   - independent upscale and hires options: run hires without upscale or upscale without hires or both
   - upscale can now run 0.1-8.0 scale and will also run if enabled at 1.0 to allow for upscalers that simply improve image quality
   - update ui section to reflect changes
   - *note*: behavior using backend:original is unchanged for backwards compatibilty
+- **Image2Video**
+  - new module for creating videos from images  
+  - simply enable from *img2img -> scripts -> image2video*  
+  - based on [VGen](https://huggingface.co/ali-vilab/i2vgen-xl)  
 - **Samplers**
   - [TCD](https://mhh0318.github.io/tcd/): Trajectory Consistency Distillation  
     new sampler that produces consistent results in a very low number of steps (comparable to LCM but without reliance on LoRA)  

modules/ui_postprocessing.py

@@ -74,16 +74,16 @@ def create_ui():
                         ui_common.create_refresh_button(clip_model, interrogate.get_clip_models, lambda: {"choices": interrogate.get_clip_models()}, 'refresh_interrogate_models')
                     with gr.Row():
                         btn_interrogate_batch = gr.Button("Interrogate", variant='primary')
-                with gr.Tab("Query Image"):
+                with gr.Tab("Visual Query"):
                     from modules import vqa
                     with gr.Row():
                         vqa_image = gr.Image(type='pil', label="Image")
                     with gr.Row():
-                        vqa_question = gr.Textbox(label="Question")
+                        vqa_question = gr.Textbox(label="Question", placeholder="Descirbe the image")
                     with gr.Row():
                         vqa_answer = gr.Textbox(label="Answer", lines=3)
                     with gr.Row():
-                        vqa_model = gr.Dropdown(list(vqa.MODELS), value='None', label='VQA Model')
+                        vqa_model = gr.Dropdown(list(vqa.MODELS), value='Moondream 2', label='VQA Model')
                         vqa_submit = gr.Button("Interrogate", variant='primary')
                     vqa_submit.click(vqa.interrogate, inputs=[vqa_question, vqa_image, vqa_model], outputs=[vqa_answer])
 

modules/vqa.py

@@ -8,7 +8,7 @@ processor = None
 model = None
 loaded: str = None
 MODELS = {
-    "None": None,
+    "Moondream 2": "vikhyatk/moondream2", # 3.7GB
     "GIT TextCaps Base": "microsoft/git-base-textcaps", # 0.7GB
     "GIT VQA Base": "microsoft/git-base-vqav2", # 0.7GB
     "GIT VQA Large": "microsoft/git-large-vqav2", # 1.6GB
@@ -40,7 +40,6 @@ def git(question: str, image: Image.Image, repo: str = None):
         generated_ids = model.generate(**git_dict)
     response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    model.to(devices.cpu)
     shared.log.debug(f'VQA: response={response}')
     return response
 
@@ -80,7 +79,6 @@ def vilt(question: str, image: Image.Image, repo: str = None):
     idx = logits.argmax(-1).item()
     response = model.config.id2label[idx]
 
-    model.to(devices.cpu)
     shared.log.debug(f'VQA: response={response}')
     return response
 
@@ -102,25 +100,51 @@ def pix(question: str, image: Image.Image, repo: str = None):
         outputs = model.generate(**inputs)
     response = processor.decode(outputs[0], skip_special_tokens=True)
 
-    model.to(devices.cpu)
     shared.log.debug(f'VQA: response={response}')
     return response
 
 
+def moondream(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = transformers.AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True) # revision = "2024-03-05"
+        processor = transformers.AutoTokenizer.from_pretrained(repo) # revision = "2024-03-05"
+        loaded = repo
+        model.eval()
+    model.to(devices.device, devices.dtype)
+    shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}')
+
+    if len(question) < 2:
+        question = "Describe the image."
+    encoded = model.encode_image(image)
+    with devices.inference_context():
+        print('HERE', question)
+        response = model.answer_question(encoded, question, processor)
+
+    shared.log.debug(f'VQA: response="{response}"')
+    return response
+
+
 def interrogate(vqa_question, vqa_image, vqa_model):
     vqa_model = MODELS.get(vqa_model, None)
-    shared.log.debug(f'VQA: model="{vqa_model}" question={vqa_question} image={vqa_image}')
+    shared.log.debug(f'VQA: model="{vqa_model}" question="{vqa_question}" image={vqa_image}')
     if vqa_image is None:
-        return 'no image provided'
+        answer = 'no image provided'
     if vqa_model is None:
-        return 'no model selected'
+        answer = 'no model selected'
     if 'git' in vqa_model.lower():
-        return git(vqa_question, vqa_image, vqa_model)
+        answer = git(vqa_question, vqa_image, vqa_model)
     if 'vilt' in vqa_model.lower():
-        return vilt(vqa_question, vqa_image, vqa_model)
+        answer = vilt(vqa_question, vqa_image, vqa_model)
     if 'blip' in vqa_model.lower():
-        return blip(vqa_question, vqa_image, vqa_model)
+        answer = blip(vqa_question, vqa_image, vqa_model)
     if 'pix' in vqa_model.lower():
-        return pix(vqa_question, vqa_image, vqa_model)
+        answer = pix(vqa_question, vqa_image, vqa_model)
+    if 'moondream2' in vqa_model.lower():
+        answer = moondream(vqa_question, vqa_image, vqa_model)
     else:
-        return 'unknown model'
+        answer = 'unknown model'
+    if model is not None:
+        model.to(devices.cpu)
+    devices.torch_gc()
+    return answer