Ketengan-Diffusion-Lab/Dolphin-Inference-MGPU

1import gradio as gr2import torch3import transformers4from transformers import AutoModelForCausalLM, AutoTokenizer5from PIL import Image6import warnings7 8# disable some warnings9transformers.logging.set_verbosity_error()10transformers.logging.disable_progress_bar()11warnings.filterwarnings('ignore')12 13model_name = 'cognitivecomputations/dolphin-vision-72b'14 15# Set up GPU memory optimization16torch.cuda.empty_cache()17device = torch.device("cuda" if torch.cuda.is_available() else "cpu")18 19# Load tokenizer20tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)21 22# Load model with memory optimizations23model = AutoModelForCausalLM.from_pretrained(24    model_name,25    torch_dtype=torch.float16,26    low_cpu_mem_usage=True,27    device_map="auto",28    trust_remote_code=True,29    offload_folder="offload",  # Offload to disk if necessary30    offload_state_dict=True,   # Offload state dict to CPU31    max_memory={0: "40GB"}     # Limit GPU memory usage32)33 34def inference(prompt, image, temperature, beam_size):35    messages = [36        {"role": "user", "content": f'<image>\n{prompt}'}37    ]38    text = tokenizer.apply_chat_template(39        messages,40        tokenize=False,41        add_generation_prompt=True42    )43 44    text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]45    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(device)46 47    image_tensor = model.process_images([image], model.config).to(device)48 49    # Clear GPU memory50    torch.cuda.empty_cache()51 52    # Generate with memory optimization53    with torch.cuda.amp.autocast():54        output_ids = model.generate(55            input_ids,56            images=image_tensor,57            max_new_tokens=1024,58            temperature=temperature,59            num_beams=beam_size,60            use_cache=True,61            do_sample=True,62            repetition_penalty=1.1,63            length_penalty=1.0,64            no_repeat_ngram_size=365        )[0]66 67    # Clear GPU memory again68    torch.cuda.empty_cache()69 70    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()71 72# Create Gradio interface73with gr.Blocks() as demo:74    with gr.Row():75        with gr.Column():76            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")77            image_input = gr.Image(label="Image", type="pil")78            temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")79            beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size")80            submit_button = gr.Button("Submit")81        with gr.Column():82            output_text = gr.Textbox(label="Output")83 84    submit_button.click(85        fn=inference, 86        inputs=[prompt_input, image_input, temperature_input, beam_size_input], 87        outputs=output_text88    )89 90# Launch the app91demo.launch()