Ketengan-Diffusion-Lab commited on
Commit
323d186
1 Parent(s): 85baff2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +57 -84
app.py CHANGED
@@ -1,7 +1,5 @@
1
- import os
2
  import gradio as gr
3
  import torch
4
- import torch.distributed as dist
5
  import transformers
6
  from transformers import AutoModelForCausalLM, AutoTokenizer
7
  from PIL import Image
@@ -12,98 +10,73 @@ transformers.logging.set_verbosity_error()
12
  transformers.logging.disable_progress_bar()
13
  warnings.filterwarnings('ignore')
14
 
15
- def setup(rank, world_size):
16
- os.environ['MASTER_ADDR'] = 'localhost'
17
- os.environ['MASTER_PORT'] = '12355'
18
- dist.init_process_group("nccl", rank=rank, world_size=world_size)
19
 
20
- def cleanup():
21
- dist.destroy_process_group()
22
 
23
- def load_model_on_gpus(model_name, num_gpus):
24
- # Calculate number of layers to assign to each GPU
25
- model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True)
26
- num_layers = len(model.model.layers)
27
- layers_per_gpu = num_layers // num_gpus
 
 
28
 
29
- # Assign layers to GPUs
30
- device_map = {}
31
- for i in range(num_layers):
32
- device_map[f'model.layers.{i}'] = i // layers_per_gpu
33
-
34
- # Assign other components
35
- device_map['model.embed_tokens'] = 0
36
- device_map['model.norm'] = num_gpus - 1
37
- device_map['lm_head'] = num_gpus - 1
38
 
39
- return AutoModelForCausalLM.from_pretrained(
40
- model_name,
41
- device_map=device_map,
42
- torch_dtype=torch.float16,
43
- trust_remote_code=True
 
 
 
44
  )
45
 
46
- def run_distributed(rank, world_size, model_name):
47
- setup(rank, world_size)
48
-
49
- if rank == 0:
50
- tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
51
-
52
- model = load_model_on_gpus(model_name, world_size)
53
-
54
- def inference(prompt, image, temperature, beam_size):
55
- if rank == 0:
56
- messages = [{"role": "user", "content": f'<image>\n{prompt}'}]
57
- text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
58
- text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
59
- input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(rank)
60
- image_tensor = model.process_images([image], model.config).to(rank)
61
- else:
62
- input_ids = torch.zeros(1, 1, dtype=torch.long).to(rank)
63
- image_tensor = torch.zeros(1, 3, 224, 224).to(rank)
64
 
65
- dist.broadcast(input_ids, src=0)
66
- dist.broadcast(image_tensor, src=0)
67
 
68
- with torch.cuda.amp.autocast():
69
- output_ids = model.generate(
70
- input_ids,
71
- images=image_tensor,
72
- max_new_tokens=1024,
73
- temperature=temperature,
74
- num_beams=beam_size,
75
- use_cache=True
76
- )[0]
77
 
78
- if rank == 0:
79
- return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
80
- else:
81
- return ""
 
 
 
 
 
 
82
 
83
- if rank == 0:
84
- with gr.Blocks() as demo:
85
- with gr.Row():
86
- with gr.Column():
87
- prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")
88
- image_input = gr.Image(label="Image", type="pil")
89
- temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
90
- beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size")
91
- submit_button = gr.Button("Submit")
92
- with gr.Column():
93
- output_text = gr.Textbox(label="Output")
94
 
95
- submit_button.click(
96
- fn=inference,
97
- inputs=[prompt_input, image_input, temperature_input, beam_size_input],
98
- outputs=output_text
99
- )
 
 
 
 
 
100
 
101
- demo.launch(share=True)
102
-
103
- cleanup()
 
 
104
 
105
- if __name__ == "__main__":
106
- model_name = 'cognitivecomputations/dolphin-vision-72b'
107
- world_size = torch.cuda.device_count()
108
- print(f"Running on {world_size} GPUs")
109
- torch.multiprocessing.spawn(run_distributed, args=(world_size, model_name), nprocs=world_size, join=True)
 
 
1
  import gradio as gr
2
  import torch
 
3
  import transformers
4
  from transformers import AutoModelForCausalLM, AutoTokenizer
5
  from PIL import Image
 
10
  transformers.logging.disable_progress_bar()
11
  warnings.filterwarnings('ignore')
12
 
13
+ # Set device to GPU if available, else CPU
14
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+ print(f"Using device: {device}")
 
16
 
17
+ model_name = 'failspy/kappa-3-phi-abliterated'
 
18
 
19
+ # create model and load it to the specified device
20
+ model = AutoModelForCausalLM.from_pretrained(
21
+ model_name,
22
+ torch_dtype=torch.float16,
23
+ device_map="auto",
24
+ trust_remote_code=True
25
+ )
26
 
27
+ tokenizer = AutoTokenizer.from_pretrained(
28
+ model_name,
29
+ trust_remote_code=True
30
+ )
 
 
 
 
 
31
 
32
+ def inference(prompt, image, temperature, beam_size):
33
+ messages = [
34
+ {"role": "user", "content": f'<image>\n{prompt}'}
35
+ ]
36
+ text = tokenizer.apply_chat_template(
37
+ messages,
38
+ tokenize=False,
39
+ add_generation_prompt=True
40
  )
41
 
42
+ text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
43
+ input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(device)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44
 
45
+ image_tensor = model.process_images([image], model.config).to(device)
 
46
 
47
+ # Add debug prints
48
+ print(f"Device of model: {next(model.parameters()).device}")
49
+ print(f"Device of input_ids: {input_ids.device}")
50
+ print(f"Device of image_tensor: {image_tensor.device}")
 
 
 
 
 
51
 
52
+ # generate
53
+ 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
+ )[0]
62
 
63
+ return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
 
 
 
 
 
 
 
 
 
 
64
 
65
+ with gr.Blocks() as demo:
66
+ with gr.Row():
67
+ with gr.Column():
68
+ prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")
69
+ image_input = gr.Image(label="Image", type="pil")
70
+ temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
71
+ beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size")
72
+ submit_button = gr.Button("Submit")
73
+ with gr.Column():
74
+ output_text = gr.Textbox(label="Output")
75
 
76
+ submit_button.click(
77
+ fn=inference,
78
+ inputs=[prompt_input, image_input, temperature_input, beam_size_input],
79
+ outputs=output_text
80
+ )
81
 
82
+ demo.launch(share=True)