Huihui-gemma-3n-E4B-it-abliterated

from transformers import AutoProcessor, Gemma3nForConditionalGeneration
from transformers.models.gemma3n.modeling_gemma3n import Gemma3nTextDecoderLayer
from PIL import Image
import requests
import torch
import jaxtyping
import einops
import base64

model_id = "google/gemma-3n-E4B-it"

model = Gemma3nForConditionalGeneration.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16,).eval()

#refusal_dir= torch.load(model_id + "/final_refusal_dirs-16.pt", map_location='cpu', weights_only=True)
#
#refusal_dir = refusal_dir.to(model.device)
#refusal_dir = refusal_dir.to(torch.bfloat16)
#
#def direction_ablation_hook(activation: jaxtyping.Float[torch.Tensor, "... d_act"],
#                            direction: jaxtyping.Float[torch.Tensor, "d_act"]):
#    proj = einops.einsum(activation, direction.view(-1, 1), '... d_act, d_act single -> ... single') * direction
#    return activation - proj
#
#class AblationDecoderLayer(Gemma3nTextDecoderLayer):
#    def __init__(self, original_layer, config, layer_idx, refusal_dir):
#        super(AblationDecoderLayer, self).__init__(config, layer_idx)
#        self.original_layer = original_layer
#        self.refusal_dir = refusal_dir
#
#    def forward(self, *args, **kwargs):
#        hidden_states = args[0]
#        ablated = direction_ablation_hook(hidden_states, self.refusal_dir.to(hidden_states.device)).to(hidden_states.device)
#        args = (ablated,) + args[1:]
#        return self.original_layer.forward(*args, **kwargs)
#
#for idx in range(len(model.model.language_model.layers)):
#    model.model.language_model.layers[idx] = AblationDecoderLayer(model.model.language_model.layers[idx], model.config.text_config, idx, refusal_dir)

processor = AutoProcessor.from_pretrained(model_id)

# https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg
image_path = model_id + "/bee.jpg"

with open(image_path, "rb") as image_file:
    encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
    image_type = "image/jpeg" if image_path.endswith(".jpg") else "image/png"
    data_uri = f"data:{image_type};base64,{encoded_string}"

    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": data_uri},
                {"type": "text", "text": "Describe this image in detail."}
            ]
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device, dtype=torch.bfloat16)

    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        generation = model.generate(**inputs, max_new_tokens=1024, do_sample=True)
        generation = generation[0][input_len:]

    decoded = processor.decode(generation, skip_special_tokens=True)
    print(decoded)


    # ## A Close-Up of a Busy Bumblebee on a Pink Cosmos Flower
    # 
    # This image is a charming, close-up shot of a vibrant pink cosmos flower, with a cute bumblebee taking center stage!
    # 
    # **Overall Composition:**
    # 
    # The photograph is well-composed, focusing primarily on a single, large cosmos flower with several other blooms scattered in the background. It has a slightly elevated perspective, making the subject feel quite prominent and engaging.
    # 
    # **The Main Subject - The Pink Cosmos:**
    # 
    # * **Color:** The star of the show is a beautiful, bright pink cosmos flower. The petals are a lovely shade, ranging from a slightly lighter to a deeper rose hue.
    # * **Shape:** The flower is roughly circular with 5 distinct petals radiating from a central disc. The petals are quite full and have a slightly pointed tip.
    # * **Details:** The central disc of the flower is a sunny yellow, adding a lovely contrast. The edges of the petals have a slight ruffled look, giving the flower some personality.
    # 
    # **The Bumblebee - A Busy Friend:**
    # 
    # * **Position:** A charming bumblebee is perched right in the middle of the pink flower, giving it a welcoming and lively feel.
    # * **Size & Color:** The bee is a classic black and yellow mix, with a fuzzy body. It's a good-sized bee, making it easily noticeable.
    # * **Pose:** The bee is positioned slightly diagonally, ensuring the viewer can see its head and legs clearly.
    # 
    # **Background Elements:**
    # 
    # * **Other Flowers:** Several other cosmos flowers are visible in the background, creating depth and visual interest. Some are slightly out of focus, while others are closer to the camera. These range in color from a lighter pink to a deeper magenta, and even a vibrant red.
    # * **Green Foliage:** Lush green leaves are visible, adding a natural backdrop to the scene. The leaves are fairly large and contribute to the overall cheerful aesthetic.
    # * **Rustic Details:** Scattered throughout are some dried flower heads, hinting at the abundance of cosmos plants in the garden.
    # 
    # **Overall Impression:**
    # 
    # The image is bright, cheerful, and full of life! The combination of the pink cosmos and the adorable bumblebee creates a delightful scene, perfect for anyone who loves gardening or insects. It's a clear and well-lit shot, making it a delightful visual treat.
    # 
    # In summary, this image captures a happy little world, where a diligent bumblebee is hard at work on a beautiful pink cosmos flower!