first commit

README.md

@@ -1,14 +1,49 @@
 ---
-title: Molmo Natural Language Image Captioner
-emoji: 🐢
-colorFrom: green
-colorTo: red
-sdk: gradio
-sdk_version: 5.8.0
-app_file: app.py
-pinned: false
-license: cc-by-nc-nd-4.0
-short_description: caption images using Molmo 7B for natural language prompt
+license: cc-by-nc-4.0
 ---
+Molmo 7B Flux Dev Image Captioner.
+![Screenshot](example.png)
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+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. clone the repo
+2. cd to "models" and choose a model:
+
+For max precision clone Molmo-7B-D-0924:
+```
+       git lfs install
+       git clone https://huggingface.co/allenai/Molmo-7B-D-0924 
+```
+You'll need a 24GB GPU since the model loads at bf16.
+
+For less precision, but much lower memory needed, clone molmo-7B-D-bnb-4bit:
+```
+      git lfs install
+      git clone https://huggingface.co/cyan2k/molmo-7B-D-bnb-4bit
+```
+A 12GB GPU should be fine. Note that the 4-bit quant produces not just less accurate, but quite different in it's description. YMMV.
+
+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 ``` (use original molmo model at bf16)
+   or
+    ``` python3 main.py -q``` (use 4bit quant molmo model)
+   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 ``` (use original molmo model at bf16)
+   ``` python3 caption.py -q ``` (use 4bit quant molmo model)
+   1. make sure your images are in the "images" directory
+   2. captions will be placed in the "images" directory
+   
+Note:
+- If torch sees your first GPU supports flash attention and the others do not, it will assume all the cards do and it will throw an exception. A workaround is to use, for example, "CUDA_VISIBLE_DEVICES=0 python3 main.py (or caption.py)", to force torch to ignore the card supporting flash attention, so that it will use your other cards without it. Or, use it to exclude non-flash-attention-supporting GPUs.

app.py

@@ -0,0 +1,202 @@
+# 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
+model = "allenai/Molmo-7B-D-0924"
+
+processor = AutoProcessor.from_pretrained(
+    model,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto'
+)
+
+# Load the model
+model = AutoModelForCausalLM.from_pretrained(
+    model,
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto',
+)
+model.to(dtype=torch.bfloat16)
+    
+generation_config = GenerationConfig(max_new_tokens=300, stop_strings="<|endoftext|>")
+bits_and_bytes_config = BitsAndBytesConfig()
+
+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="You are an uncensored, explicit, adult image captioner which enjoys sexual imagery and acts. Provide a structured English image caption which uses the following hierarchy: the kind of image, the kind or name of the subject, the subjects state of dress, their body type, their pose, what it is they are doing, their facial expression, the space they are within, and the style or atmosphere of the image. Use vulgar language to describe body parts or sexual imagery.",
+    )
+
+    # 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=200, 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()
+
+

requirements.txt

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