Ketengan-Diffusion-Lab commited on
Commit
9aeab55
·
verified ·
1 Parent(s): f81e89d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -3
app.py CHANGED
@@ -20,7 +20,7 @@ model_name = 'cognitivecomputations/dolphin-vision-72b'
20
  model = AutoModelForCausalLM.from_pretrained(
21
  model_name,
22
  torch_dtype=torch.float16,
23
- device_map="auto", # This will automatically use the GPU if available
24
  trust_remote_code=True
25
  )
26
 
@@ -29,7 +29,7 @@ tokenizer = AutoTokenizer.from_pretrained(
29
  trust_remote_code=True
30
  )
31
 
32
- def inference(prompt, image):
33
  messages = [
34
  {"role": "user", "content": f'<image>\n{prompt}'}
35
  ]
@@ -55,6 +55,8 @@ def inference(prompt, image):
55
  input_ids,
56
  images=image_tensor,
57
  max_new_tokens=1024,
 
 
58
  use_cache=True
59
  )[0]
60
 
@@ -65,10 +67,16 @@ with gr.Blocks() as demo:
65
  with gr.Column():
66
  prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")
67
  image_input = gr.Image(label="Image", type="pil")
 
 
68
  submit_button = gr.Button("Submit")
69
  with gr.Column():
70
  output_text = gr.Textbox(label="Output")
71
 
72
- submit_button.click(fn=inference, inputs=[prompt_input, image_input], outputs=output_text)
 
 
 
 
73
 
74
  demo.launch(share=True)
 
20
  model = AutoModelForCausalLM.from_pretrained(
21
  model_name,
22
  torch_dtype=torch.float16,
23
+ device_map="auto",
24
  trust_remote_code=True
25
  )
26
 
 
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
  ]
 
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
 
 
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)