Update app.py

app.py

@@ -1,7 +1,5 @@
-import os
 import gradio as gr
 import torch
-import torch.distributed as dist
 import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
@@ -12,98 +10,73 @@ 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)
+# Set device to GPU if available, else CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
-def cleanup():
-    dist.destroy_process_group()
+model_name = 'failspy/kappa-3-phi-abliterated'
 
-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
+# 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
+)
 
-    # 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
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True
+)
 
-    return AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map=device_map,
-        torch_dtype=torch.float16,
-        trust_remote_code=True
+def inference(prompt, image, temperature, beam_size):
+    messages = [
+        {"role": "user", "content": f'<image>\n{prompt}'}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=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)
+    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)
 
-        dist.broadcast(input_ids, src=0)
-        dist.broadcast(image_tensor, src=0)
+    image_tensor = model.process_images([image], model.config).to(device)
 
-        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]
+    # 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}")
 
-        if rank == 0:
-            return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
-        else:
-            return ""
+    # generate
+    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:
-        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")
+    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
 
-            submit_button.click(
-                fn=inference, 
-                inputs=[prompt_input, image_input, temperature_input, beam_size_input], 
-                outputs=output_text
-            )
+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")
 
-        demo.launch(share=True)
-
-    cleanup()
+    submit_button.click(
+        fn=inference, 
+        inputs=[prompt_input, image_input, temperature_input, beam_size_input], 
+        outputs=output_text
+    )
 
-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)
+demo.launch(share=True)