import os

from transformers import AutoModelForMultimodalLM, AutoProcessor, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]), # Number of devices across which the dense layers are sharded
    ep_size=int(os.environ["WORLD_SIZE"]), # Number of devices across which the experts are distributed
)

processor = AutoProcessor.from_pretrained("meta-models/Muse-Spark-Open-816B-A42B")
model = AutoModelForMultimodalLM.from_pretrained(
    "meta-models/Muse-Spark-Open-816B-A42B",
    distributed_config=distributed_config,
)

messages = [
    {"role": "user", "content": "Explain in one sentence why the sky is usually blue."},
]
inputs = processor.apply_chat_template(
    messages,
    reasoning_effort="low",  # default: xhigh, options: low/medium/high/xhigh
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=512)
print(processor.decode(outputs[0][input_len:], skip_special_tokens=False))
#  to=self<|message|>Explain in one sentence why the sky is usually blue.
# ...
# <|eom|><|start|>assistant to=user<|message|>The sky is usually blue because molecules in Earth's atmosphere scatter shorter-wavelength blue sunlight in all directions much more strongly than longer-wavelength red light.<|eot|>

import os

from transformers import AutoModelForMultimodalLM, AutoProcessor, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

processor = AutoProcessor.from_pretrained("meta-models/Muse-Spark-Open-816B-A42B")
model = AutoModelForMultimodalLM.from_pretrained(
    "meta-models/Muse-Spark-Open-816B-A42B",
    distributed_config=distributed_config,
)

messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
            {"type": "text", "text": "What is shown in this image?"},
        ],
    },
]
inputs = processor.apply_chat_template(
    messages,
    reasoning_effort="medium", # default: xhigh, options: low/medium/high/xhigh
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=512)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)
print(response)
#  to=self<|message|>What is shown in this image?
# ...
# <|eom|><|start|>assistant to=user<|message|>The image shows a **Pallas's cat — also known as a manul, *Otocolobus manul*** — walking through snow.
#
# It is a small, stocky wild cat recognizable by:
#
# * **Its extremely dense, long, greyish-brown winter coat** with faint dark stripes and a dusting of snow on its back
# * **Its flat, broad face** with low-set ears, a short muzzle, and dark stripes on the cheeks
# * **Its short legs and thick, plump body** adapted for cold climates
#
# The cat is captured mid-stride on a snow-covered surface, with one front paw lifted. In the background are white birch tree trunks with black markings and a chain-link fence on the left, indicating the animal is in a zoo or wildlife enclosure in winter.<|eot|>

import os

from transformers import AutoModelForMultimodalLM, AutoProcessor, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

processor = AutoProcessor.from_pretrained("meta-models/Muse-Spark-Open-816B-A42B")
model = AutoModelForMultimodalLM.from_pretrained(
    "meta-models/Muse-Spark-Open-816B-A42B",
    distributed_config=distributed_config,
)

messages = [
    {
        "role": "user",
        "content": [
            {"type": "video", "url": "https://huggingface.co/datasets/hf-internal-testing/fixtures_videos/resolve/main/tennis.mp4"},
            {"type": "text", "text": "What is happening in this video?"},
        ],
    },
]
inputs = processor.apply_chat_template(
    messages,
    reasoning_effort="medium",  # default: xhigh, options: low/medium/high/xhigh
    processor_kwargs={
        "fps": 2,   # default: 2, number of frames sampled per second
        "num_frames": 8, # default: 94, maximum number of frames sampled from the video
    },
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=512)
print(processor.decode(outputs[0][input_len:], skip_special_tokens=False))
#  to=self<|message|>What is happening in this video?
# ...
# <|eom|><|start|>assistant to=user<|message|>In this video two people are playing tennis on an indoor court.
#
# The shot is taken from behind one of the players - a man in a white t-shirt, black shorts and white shoes and socks. He is standing near the baseline on an orange hard court with white lines, in a ready position with his knees bent and his racket held in front of him.
#
# Across the net, on the far side of the court, is his opponent, also dressed in light-colored clothing.
#
# The sequence shows a short rally:
#
# *   The opponent hits the ball over the net.
# *   The man in the foreground shuffles to his left, turns his body, and prepares for a shot.
# *   He executes a two-handed backhand, swinging the racket across his body and following through as he hits the ball back over the net.
#
# The setting is a large indoor tennis hall with a high, arched ceiling with exposed wooden beams and overhead fluorescent lights. The court is surrounded by a grey run-off area, green windscreening on the walls, and netting that separates it from adjacent courts. On the right side you can see an umpire's chair and a few yellow and black benches. An "Adobe Stock" watermark is visible in the center of the frame.<|eot|>

import os

from transformers import AutoModelForMultimodalLM, AutoProcessor, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

processor = AutoProcessor.from_pretrained("meta-models/Muse-Spark-Open-816B-A42B")
model = AutoModelForMultimodalLM.from_pretrained(
    "meta-models/Muse-Spark-Open-816B-A42B",
    distributed_config=distributed_config,
)

messages = [
    {
        "role": "user",
        "content": [
            {"type": "audio", "url": "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/glass_breaking.mp3"},
            {"type": "text", "text": "What is said in this recording?"},
        ],
    },
]
inputs = processor.apply_chat_template(
    messages,
    reasoning_effort="low", # default: xhigh, options: low/medium/high/xhigh
    processor_kwargs={"sampling_rate": 24000},  # sampling rate the model was trained with
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=512)
print(processor.decode(outputs[0][input_len:], skip_special_tokens=False))
#  to=self<|message|>We need to respond to user request: "What is said in this recording?"
# ...
# <|eom|><|start|>assistant to=user<|message|>In the recording the person says:
#
# > "Just picked up this glass and the people at the store told me that if I dropped it it wouldn't break."
#
# Then the glass is dropped and shatters.<|eot|>