Update app.py

app.py

@@ -4,7 +4,8 @@ import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 import warnings
-from accelerate import Accelerator
+from accelerate import Accelerator, DistributedType
+import os
 
 # disable some warnings
 transformers.logging.set_verbosity_error()
@@ -16,19 +17,22 @@ accelerator = Accelerator()
 
 model_name = 'cognitivecomputations/dolphin-vision-72b'
 
-# 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
-    )
+# Determine the number of GPUs available
+num_gpus = torch.cuda.device_count()
+print(f"Number of GPUs available: {num_gpus}")
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        trust_remote_code=True
-    )
+# Load model and tokenizer
+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
+)
 
 # Prepare model
 model = accelerator.prepare(model)
@@ -63,31 +67,25 @@ 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()
 
-# 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)
-
-# Wait for all processes to finish
-accelerator.wait_for_everyone()
+# Create Gradio interface
+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
+    )
+
+# Launch the app
+demo.launch()