Create app.py

app.py

@@ -0,0 +1,72 @@
+import gradio as gr
+from transformers import (
+    AutoModelForCausalLM,
+    AutoProcessor,
+    GenerationConfig,
+    BitsAndBytesConfig,
+)
+from PIL import Image
+import torch
+
+# Configuration for 4-bit quantization and GPU offloading
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+)
+
+# Model repository
+repo_name = "cyan2k/molmo-7B-O-bnb-4bit"
+
+# Load the processor and model
+processor = AutoProcessor.from_pretrained(repo_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True,
+    quantization_config=bnb_config,
+)
+
+# Ensure model is on GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+def describe_images(images):
+    descriptions = []
+    for image in images:
+        if isinstance(image, str):
+            image = Image.open(image)
+        # Process the image
+        inputs = processor.process(
+            images=[image],
+            text="Describe this image in great detail.",
+        )
+        # Move inputs to the same device as the model
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        # Generate output
+        with torch.no_grad():
+            output = model.generate_from_batch(
+                inputs,
+                GenerationConfig(max_new_tokens=200, stop_strings=["<|endoftext|>"]),
+                tokenizer=processor.tokenizer,
+            )
+        # Decode generated tokens to text
+        generated_tokens = output[0, inputs["input_ids"].size(1):]
+        generated_text = processor.tokenizer.decode(
+            generated_tokens, skip_special_tokens=True
+        )
+        descriptions.append(generated_text.strip())
+    return "\n\n".join(descriptions)
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("<h3><center>Image Description Generator</center></h3>")
+    with gr.Row():
+        image_input = gr.File(
+            file_types=["image"], label="Upload Image(s)", multiple=True
+        )
+    generate_button = gr.Button("Generate Descriptions")
+    output_text = gr.Textbox(label="Descriptions", lines=15)
+
+    generate_button.click(describe_images, inputs=image_input, outputs=output_text)
+
+demo.launch()