sandz7 commited on
Commit
5ae7f9c
β€’
1 Parent(s): 153de5a

added ** inside steamer

Browse files
Files changed (1) hide show
  1. app.py +19 -8
app.py CHANGED
@@ -5,6 +5,7 @@ from PIL import Image
5
  import threading
6
  import spaces
7
  import accelerate
 
8
 
9
  DESCRIPTION = '''
10
  <div>
@@ -56,14 +57,14 @@ def krypton(input,
56
  # Image is not defined at all
57
  gr.Error("Uplaod an image for Krypton to work")
58
 
59
- image = Image.open(image)
60
- # image = Image.open(requests.get(url, stream=True).raw)
61
  prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\n{input['text']}<|eot_id|>"
62
  "<|start_header_id|>assistant<|end_header_id|>\n\n")
63
- inputs = processor(prompt, image, return_tensors='pt').to('cuda', torch.float16)
 
 
64
 
65
  # Streamer
66
- streamer = TextIteratorStreamer(processor, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
67
 
68
  if temperature == 0.0:
69
  do_sample = False
@@ -81,10 +82,20 @@ def krypton(input,
81
  thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
82
  thread.start()
83
 
84
- outputs = []
85
- for text in streamer:
86
- outputs.append(text)
87
- yield "".join(outputs)
 
 
 
 
 
 
 
 
 
 
88
 
89
  chatbot=gr.Chatbot(height=600, label="Krypt AI")
90
  chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter your question or upload an image.", show_label=False)
 
5
  import threading
6
  import spaces
7
  import accelerate
8
+ import time
9
 
10
  DESCRIPTION = '''
11
  <div>
 
57
  # Image is not defined at all
58
  gr.Error("Uplaod an image for Krypton to work")
59
 
 
 
60
  prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\n{input['text']}<|eot_id|>"
61
  "<|start_header_id|>assistant<|end_header_id|>\n\n")
62
+
63
+ image = Image.open(image)
64
+ inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
65
 
66
  # Streamer
67
+ streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
68
 
69
  if temperature == 0.0:
70
  do_sample = False
 
82
  thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
83
  thread.start()
84
 
85
+ buffer = ""
86
+ time.sleep(0.5)
87
+ for new_text in streamer:
88
+ # find <|eot_id|> and remove it from the new_text
89
+ if "<|eot_id|>" in new_text:
90
+ new_text = new_text.split("<|eot_id|>")[0]
91
+ buffer += new_text
92
+
93
+ # generated_text_without_prompt = buffer[len(text_prompt):]
94
+ generated_text_without_prompt = buffer
95
+ # print(generated_text_without_prompt)
96
+ time.sleep(0.06)
97
+ # print(f"new_text: {generated_text_without_prompt}")
98
+ yield generated_text_without_prompt
99
 
100
  chatbot=gr.Chatbot(height=600, label="Krypt AI")
101
  chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter your question or upload an image.", show_label=False)