Spaces:
Build error
Build error
| import os | |
| import gradio as gr | |
| import torch | |
| import torch.distributed as dist | |
| import transformers | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from PIL import Image | |
| import warnings | |
| # disable some warnings | |
| transformers.logging.set_verbosity_error() | |
| transformers.logging.disable_progress_bar() | |
| warnings.filterwarnings('ignore') | |
| def setup(rank, world_size): | |
| os.environ['MASTER_ADDR'] = 'localhost' | |
| os.environ['MASTER_PORT'] = '12355' | |
| dist.init_process_group("nccl", rank=rank, world_size=world_size) | |
| def cleanup(): | |
| dist.destroy_process_group() | |
| def load_model_on_gpus(model_name, num_gpus): | |
| # Calculate number of layers to assign to each GPU | |
| model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True) | |
| num_layers = len(model.model.layers) | |
| layers_per_gpu = num_layers // num_gpus | |
| # Assign layers to GPUs | |
| device_map = {} | |
| for i in range(num_layers): | |
| device_map[f'model.layers.{i}'] = i // layers_per_gpu | |
| # Assign other components | |
| device_map['model.embed_tokens'] = 0 | |
| device_map['model.norm'] = num_gpus - 1 | |
| device_map['lm_head'] = num_gpus - 1 | |
| return AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| device_map=device_map, | |
| torch_dtype=torch.float16, | |
| trust_remote_code=True | |
| ) | |
| def run_distributed(rank, world_size, model_name): | |
| setup(rank, world_size) | |
| if rank == 0: | |
| tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) | |
| model = load_model_on_gpus(model_name, world_size) | |
| def inference(prompt, image, temperature, beam_size): | |
| if rank == 0: | |
| messages = [{"role": "user", "content": f'<image>\n{prompt}'}] | |
| text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')] | |
| input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(rank) | |
| image_tensor = model.process_images([image], model.config).to(rank) | |
| else: | |
| input_ids = torch.zeros(1, 1, dtype=torch.long).to(rank) | |
| image_tensor = torch.zeros(1, 3, 224, 224).to(rank) | |
| dist.broadcast(input_ids, src=0) | |
| dist.broadcast(image_tensor, src=0) | |
| with torch.cuda.amp.autocast(): | |
| output_ids = model.generate( | |
| input_ids, | |
| images=image_tensor, | |
| max_new_tokens=1024, | |
| temperature=temperature, | |
| num_beams=beam_size, | |
| use_cache=True | |
| )[0] | |
| if rank == 0: | |
| return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip() | |
| else: | |
| return "" | |
| if rank == 0: | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail") | |
| image_input = gr.Image(label="Image", type="pil") | |
| temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature") | |
| beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size") | |
| submit_button = gr.Button("Submit") | |
| with gr.Column(): | |
| output_text = gr.Textbox(label="Output") | |
| submit_button.click( | |
| fn=inference, | |
| inputs=[prompt_input, image_input, temperature_input, beam_size_input], | |
| outputs=output_text | |
| ) | |
| demo.launch(share=True) | |
| cleanup() | |
| if __name__ == "__main__": | |
| model_name = 'cognitivecomputations/dolphin-vision-72b' | |
| world_size = torch.cuda.device_count() | |
| print(f"Running on {world_size} GPUs") | |
| torch.multiprocessing.spawn(run_distributed, args=(world_size, model_name), nprocs=world_size, join=True) |