added processor and conditional generation

app.py

@@ -1,9 +1,9 @@
-# import torch
-# import gradio as gr
-# from transformers import pipeline, TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
-# from PIL import Image
-# import requests
-# import threading
+import torch
+import gradio as gr
+from transformers import TextIteratorStreamer, AutoProcessor, LlavaForConditionalGeneration
+from PIL import Image
+import requests
+import threading
 
 DESCRIPTION = '''
 <div>
@@ -12,63 +12,14 @@ DESCRIPTION = '''
 </div>
 '''
 
-# model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
-# pipe = pipeline("image-to-text", model=model_id, device_map="auto")
-# # Place transformers in hardware to prepare for process and generation
-# llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
-# llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", torch_dtype=torch.float16).to('cuda')
-# terminators = [
-#     llama_tokenizer.eos_token_id,
-#     llama_tokenizer.convert_tokens_to_ids("<|eot_id|>")
-# ]
-
-# def krypton(prompt,
-#             history,
-#             input_image,
-#             max_new_tokens,
-#             temperature,
-#             num_beams,
-#             do_sample: bool=True):
-#     """
-#     Passes an image as input, places it for generation
-#     on pipeline and output is passed. This is multimodal
-#     """
-#     conversation = []
-#     for user, assistant in history:
-#         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-#     conversation.append({"role": "user", "content": prompt})
-
-#     input_ids = llama_tokenizer.apply_chat_template(conversation, return_tensors='pt').to(llama_model.device)
-
-#     streamer = TextIteratorStreamer(llama_tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-
-#     llava_generation_kwargs = dict(
-#         input_ids=input_ids,
-#         streamer=streamer,
-#         max_new_tokens=max_new_tokens,
-#         num_beams=num_beams,
-#         do_sample=do_sample
-#     )
-
-#     if temperature == 0.0:
-#         do_sample = False
-
-#     pil_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
-
-#     # Pipeline generation
-#     outputs = pipeline()
-
-
-from transformers import pipeline
-from PIL import Image    
-import requests
-import torch
-import subprocess
-import gradio as gr
-import spaces
-
 model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
-pipe = pipeline("image-to-text", model=model_id, torch_dtype=torch.float16, device=0)
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True
+).to('cuda')
+
+processor = AutoProcessor.from_pretrained(model_id)
 
 @spaces.GPU(duration=120)
 def krypton(input_image):
@@ -77,9 +28,10 @@ def krypton(input_image):
     # image = Image.open(requests.get(url, stream=True).raw)
     prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\nWhat are these?<|eot_id|>"
             "<|start_header_id|>assistant<|end_header_id|>\n\n")
-    outputs = pipe(pil_image, prompt=prompt, generate_kwargs={"max_new_tokens": 200})
-    nvidia_result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE)
-    return outputs[0]
+    inputs = processor(prompt, pil_image, return_tensors='pt').to('cuda', torch.float16)
+    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False)
+    output_text = processor.decode(outputs[0][:2], skip_special_tokens=True)
+    return output_text
 
 with gr.Blocks(fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)