first commit

README.md

@@ -1,3 +1,26 @@
----
-license: cc-by-nc-4.0
----
+---
+license: cc-by-nc-4.0
+---
+Molmo 7B Flux Dev Image Captioner.
+![Screenshot](example.png)
+
+A simple python and gradio script to use Molmo 7B for image captioning. The prompt is currently written to produce captions that work well for Flux Dev LoRA training, but you could adjust it to suit other models captioning style.
+
+Install:
+1. create a python3 venv or use conda to create an environment, eg:
+   ``` conda create -n caption python=3.11 ```
+2. activate your environment, eg:
+   ``` conda activate caption ```
+3. install the dependencies
+   ``` pip3 install -r requirements.txt ```
+4. run the gradio version:
+   ``` python3 main.py ```
+   1. create a zip file of images
+   2. upload it
+   3. process it
+   4. click the button to download the caption zip file, the link is at the top of the page
+
+   run the command-line version:
+   ``` python3 caption.py ```
+   1. make sure your images are in the "images" directory
+   2. captions will be placed in the "images" directory

caption.py

@@ -0,0 +1,74 @@
+import os
+import torch
+from PIL import Image
+import requests
+from transformers import AutoProcessor, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print("GPU is available. Using CUDA.")
+else:
+    device = torch.device("cpu")
+    print("GPU is not available. Using CPU.")
+
+# load the processor from local path
+local_path = "./model/Molmo-7B-D-0924"
+processor = AutoProcessor.from_pretrained(
+    local_path,
+    local_files_only=True,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto'
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    local_path,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto',
+)    
+
+
+model.to(dtype=torch.bfloat16)
+
+# directory containing the images
+image_directory = "./images"
+
+# iterate through the images in the directory
+for filename in os.listdir(image_directory):
+    if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):  # add more image extensions if needed
+        image_path = os.path.join(image_directory, filename)
+        image = Image.open(image_path)
+        
+        # process the image and text
+        inputs = processor.process(
+            images=[image],
+            text="Describe what you see in vivid detail, without line breaks. Include information about the pose of characters, their facial expression, their height, body type, weight, the position of their limbs, and the direction of their gaze, the color of their eyes, hair, and skin. If you know a person or place name, provide it. If you know the name of an artist who may have created what you see, provide that. Do not provide opinions or value judgements. Limit your response to 276 words to avoid your description getting cut off.",
+        )
+
+        # move inputs to the correct device and make a batch of size 1
+        inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+        inputs["images"] = inputs["images"].to(torch.bfloat16)
+
+        # generate output; maximum 500 new tokens; stop generation when   is generated
+        with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
+            output = model.generate_from_batch(
+                inputs,
+                GenerationConfig(max_new_tokens=500, stop_strings="<|endoftext|>"),
+                tokenizer=processor.tokenizer,
+            )
+
+        # only get generated tokens; decode them to text
+        generated_tokens = output[0, inputs["input_ids"].size(1) :]
+        generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+        # print the generated text
+        print("Caption for: ", filename)
+        print(generated_text)
+        # print a divider
+        print("*---------------------------------------------------*")
+
+        # save the generated text to a file
+        output_filename = os.path.splitext(filename)[0] + ".txt"
+        with open(os.path.join(image_directory,output_filename), "w") as file:
+            file.write(generated_text)

main.py

@@ -0,0 +1,207 @@
+# note: if you have a mix of Ampere and newer, and also older than Ampere GPUs, set the environment variable
+#    CUDA_VISIBLE_DEVICE=1,2,3 (for example) so that one or the other is excluded.
+#    otherwise the script may fail with a flash attention exception.
+
+import gradio as gr
+import os
+import uuid
+import zipfile
+import torch
+from PIL import Image
+import requests
+from transformers import AutoProcessor, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
+from io import BytesIO
+import base64
+import atexit
+import shutil
+
+
+def cleanup_temp_files():
+    # Delete the subdirectories inside the "images" directory
+    if os.path.exists("images"):
+        for dir_name in os.listdir("images"):
+            dir_path = os.path.join("images", dir_name)
+            if os.path.isdir(dir_path):
+                shutil.rmtree(dir_path)
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print("GPU is available. Using CUDA.")
+else:
+    device = torch.device("cpu")
+    print("GPU is not available. Using CPU.")
+
+# load the processor from local path
+local_path = "./model/Molmo-7B-D-0924"
+#print("Loading processor from local path...")
+processor = AutoProcessor.from_pretrained(
+    local_path,
+    local_files_only=True,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto'
+)
+#print("Processor loaded.")
+
+print("Loading model from local path...")
+model = AutoModelForCausalLM.from_pretrained(
+    local_path,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto',
+)
+#print("Model loaded.")
+
+generation_config = GenerationConfig(max_new_tokens=300, stop_strings="<|endoftext|>")
+bits_and_bytes_config = BitsAndBytesConfig()
+
+# load the model in bf16 to reduce VRAM needed
+model.to(dtype=torch.bfloat16)
+#print("Model loaded in bf16")
+
+def unzip_images(zip_file):
+    # Create a unique directory for extracted images inside the "images" directory
+    session_dir = os.path.join("images", str(uuid.uuid4()))
+    os.makedirs(session_dir, exist_ok=True)
+    
+    # Extract images from the ZIP file to the session directory
+    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+        for file_info in zip_ref.infolist():
+            if not file_info.is_dir() and not file_info.filename.startswith("__MACOSX") and not file_info.filename.startswith("."):
+                zip_ref.extract(file_info, session_dir)
+    
+    # Get the list of image paths
+    image_paths = [os.path.join(session_dir, filename) for filename in os.listdir(session_dir) if filename.lower().endswith(('.jpg', '.jpeg', '.png'))]
+    
+    # Read the image data as PIL Image objects for previews
+    image_data = []
+    for image_path in image_paths:
+        image = Image.open(image_path)
+        image.thumbnail((128, 128))  # Resize the image to a maximum size of 128x128 pixels
+        image_data.append(image)
+    
+    # Return the list of image paths and resized image data for previews
+    return image_paths, image_data
+
+def generate_caption(image_path, processor, model, generation_config, bits_and_bytes_config):
+    # generate a caption and return it
+    caption = f"Caption for {image_path}"
+    
+    print("Processing ", image_path)
+    
+    image = Image.open(image_path)
+    # process the image and text
+    inputs = processor.process(
+        images=[image],
+        text="Describe what you see in vivid detail, without line breaks. Include information about the pose of characters, their facial expression, their height, body type, weight, the position of their limbs, and the direction of their gaze, the color of their eyes, hair, and skin. If you know a person or place name, provide it. If you know the name of an artist who may have created what you see, provide that. Do not provide opinions or value judgements. Limit your response to 276 words to avoid your description getting cut off.",
+    )
+
+    # move inputs to the correct device and make a batch of size 1
+    inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+    inputs["images"] = inputs["images"].to(torch.bfloat16)
+
+    # generate output; maximum 500 new tokens; stop generation when   is generated
+    with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
+        output = model.generate_from_batch(
+            inputs,
+            GenerationConfig(max_new_tokens=500, stop_strings="<|endoftext|>"),
+            tokenizer=processor.tokenizer,
+        )
+
+    # only get generated tokens; decode them to text
+    generated_tokens = output[0, inputs["input_ids"].size(1) :]
+    generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+    # return the generated text
+    return generated_text
+
+def process_images(image_paths, image_data):
+    captions = []
+    session_dir = os.path.dirname(image_paths[0])
+    
+    for image_path in image_paths:
+        filename = os.path.basename(image_path)  # Add this line to get the filename
+        if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
+            # Process the image using the loaded model
+            # Use the loaded model to generate the caption
+            caption = generate_caption(image_path, processor, model, generation_config, bits_and_bytes_config)
+            captions.append(caption)
+            
+            # Save the caption to a text file
+            with open(os.path.join(session_dir, f"{os.path.splitext(filename)[0]}.txt"), 'w') as f:
+                f.write(caption)
+    
+    # Create a ZIP file containing the caption text files
+    zip_filename = f"{session_dir}.zip"
+    with zipfile.ZipFile(zip_filename, 'w') as zip_ref:
+        for filename in os.listdir(session_dir):
+            if filename.lower().endswith('.txt'):
+                zip_ref.write(os.path.join(session_dir, filename), filename)
+    
+    # Delete the session directory and its contents
+    for filename in os.listdir(session_dir):
+        os.remove(os.path.join(session_dir, filename))
+    os.rmdir(session_dir)
+    
+    return captions, zip_filename, image_paths
+    
+def format_captioned_image(image, caption):
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    return f"<img src='data:image/jpeg;base64,{encoded_image}' style='width: 128px; height: 128px; object-fit: cover; margin-right: 8px;' /><span>{caption}</span>"
+
+def process_images_and_update_gallery(zip_file):
+    image_paths, image_data = unzip_images(zip_file)
+    captions, zip_filename, image_paths = process_images(image_paths, image_data)
+    image_captions = [format_captioned_image(img, caption) for img, caption in zip(image_data, captions)]
+    return gr.Markdown("\n".join(image_captions)), zip_filename
+
+def main():
+    # Register the cleanup function to be called on program exit
+    atexit.register(cleanup_temp_files)
+
+    with gr.Blocks(css="""
+        .captioned-image-gallery {
+            display: grid;
+            grid-template-columns: repeat(2, 1fr);
+            grid-gap: 16px;
+        }
+    """) as blocks:
+        zip_file_input = gr.File(label="Upload ZIP file containing images")
+        image_gallery = gr.Markdown(label="Image Previews")
+        submit_button = gr.Button("Submit")
+        zip_download_button = gr.Button("Download Caption ZIP", visible=False)
+        zip_filename = gr.State("")
+
+        zip_file_input.upload(
+            lambda zip_file: "\n".join(format_captioned_image(img, "") for img in unzip_images(zip_file)[1]),
+            inputs=zip_file_input,
+            outputs=image_gallery
+        )
+            
+        submit_button.click(
+            process_images_and_update_gallery,
+            inputs=[zip_file_input],
+            outputs=[image_gallery, zip_filename]
+        )
+
+        zip_filename.change(
+            lambda zip_filename: gr.update(visible=True),
+            inputs=zip_filename,
+            outputs=zip_download_button
+        )
+
+        zip_download_button.click(
+            lambda zip_filename: (gr.update(value=zip_filename), gr.update(visible=True), cleanup_temp_files()),
+            inputs=zip_filename,
+            outputs=[zip_file_input, zip_download_button]
+        )
+
+    blocks.launch(server_name='0.0.0.0')
+
+if __name__ == "__main__":
+    main()
+
+

model/Molmo-7B-D-0924

@@ -0,0 +1 @@
+Subproject commit 90426556d5eb7c123eb4368dd1768e8e77f624af

requirements.txt

@@ -0,0 +1,11 @@
+gradio
+uuid
+bitsandbytes
+accelerate
+transformers
+torch
+torchvision
+Pillow
+requests
+einops
+flash-attn

test-images.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a84383fcb27d0be0006744b76c97e77d2a45e852d4f17ae29eff2f8346b4923
+size 3069789