| import gradio as gr | |
| import torch | |
| from transformers import AutoProcessor, LlavaForConditionalGeneration | |
| from transformers import BitsAndBytesConfig | |
| from sentence_transformers import SentenceTransformer, util | |
| quantization_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.float16 | |
| ) | |
| embedder = SentenceTransformer('all-mpnet-base-v2') | |
| model_id = "llava-hf/llava-1.5-7b-hf" | |
| processor = AutoProcessor.from_pretrained(model_id) | |
| model = LlavaForConditionalGeneration.from_pretrained( | |
| model_id, | |
| quantization_config=quantization_config, | |
| device_map="auto", | |
| use_flash_attention_2=True, | |
| low_cpu_mem_usage=True | |
| ) | |
| def text_to_image(image, prompt): | |
| prompt = f'USER: <image>\n{prompt}\nASSISTANT:' | |
| inputs = processor([prompt], images=[image], padding=True, return_tensors="pt").to(model.device) | |
| output = model.generate(**inputs, max_new_tokens=500) | |
| generated_text = processor.batch_decode(output, skip_special_tokens=True) | |
| text = generated_text.pop() | |
| text_output = text.split("ASSISTANT:")[-1] | |
| text_embeddings = embedder.encode(text_output) | |
| return text_output, dict(text_embeddings=text_embeddings) | |
| demo = gr.Interface( | |
| fn=text_to_image, | |
| inputs=[ | |
| gr.Image(label='Select an image to analyze', type='pil'), | |
| gr.Textbox(label='Enter Prompt') | |
| ], | |
| outputs=[gr.Textbox(label='Maurice says:'), gr.JSON(label='Embedded text')] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(show_api=False) | |