Update app.py

app.py

@@ -4,30 +4,34 @@ import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 import warnings
+from accelerate import Accelerator
 
 # disable some warnings
 transformers.logging.set_verbosity_error()
 transformers.logging.disable_progress_bar()
 warnings.filterwarnings('ignore')
 
-# Set device to GPU if available, else CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+# Initialize Accelerator
+accelerator = Accelerator()
 
 model_name = 'cognitivecomputations/dolphin-vision-72b'
 
-# create model and load it to the specified device
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True
-)
+# Load model and tokenizer within main_process_first context
+with accelerator.main_process_first():
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True
+    )
 
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
+# Prepare model
+model = accelerator.prepare(model)
 
 def inference(prompt, image, temperature, beam_size):
     messages = [
@@ -40,18 +44,17 @@ def inference(prompt, image, temperature, beam_size):
     )
 
     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(device)
+    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
 
-    image_tensor = model.process_images([image], model.config).to(device)
+    image_tensor = model.process_images([image], model.config)
 
-    # Add debug prints
-    print(f"Device of model: {next(model.parameters()).device}")
-    print(f"Device of input_ids: {input_ids.device}")
-    print(f"Device of image_tensor: {image_tensor.device}")
+    # Move tensors to the appropriate device
+    input_ids = input_ids.to(accelerator.device)
+    image_tensor = image_tensor.to(accelerator.device)
 
     # generate
     with torch.cuda.amp.autocast():
-        output_ids = model.generate(
+        output_ids = accelerator.unwrap_model(model).generate(
             input_ids,
             images=image_tensor,
             max_new_tokens=1024,
@@ -60,23 +63,31 @@ def inference(prompt, image, temperature, beam_size):
             use_cache=True
         )[0]
 
+    # Gather output from all processes
+    output_ids = accelerator.gather(output_ids)
+
     return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
 
-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
-    )
+# Only create and launch Gradio interface on the main process
+if accelerator.is_main_process:
+    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)
 
-demo.launch(share=True)
+# Wait for all processes to finish
+accelerator.wait_for_everyone()