Spaces:
Build error
Build error
| from transformers import ViTConfig, ViTForImageClassification | |
| from transformers import ViTFeatureExtractor | |
| from PIL import Image | |
| import requests | |
| import matplotlib.pyplot as plt | |
| import gradio as gr | |
| from gradio.mix import Parallel | |
| from transformers import ImageClassificationPipeline, PerceiverForImageClassificationConvProcessing, PerceiverFeatureExtractor | |
| from transformers import VisionEncoderDecoderModel | |
| from transformers import AutoTokenizer | |
| import torch | |
| from transformers import ( | |
| AutoModelForCausalLM, | |
| LogitsProcessorList, | |
| MinLengthLogitsProcessor, | |
| StoppingCriteriaList, | |
| MaxLengthCriteria, | |
| ) | |
| import json | |
| import os | |
| from screenshot import ( | |
| before_prompt, | |
| prompt_to_generation, | |
| after_generation, | |
| js_save, | |
| js_load_script, | |
| ) | |
| from spaces_info import description, examples, initial_prompt_value | |
| API_URL = os.getenv("API_URL") | |
| HF_API_TOKEN = os.getenv("HF_API_TOKEN") | |
| def query(payload): | |
| print(payload) | |
| response = requests.request("POST", API_URL, json=payload, headers={"Authorization": f"Bearer {HF_API_TOKEN}"}) | |
| print(response) | |
| return json.loads(response.content.decode("utf-8")) | |
| # https://github.com/NielsRogge/Transformers-Tutorials/blob/master/HuggingFace_vision_ecosystem_overview_(June_2022).ipynb | |
| # option 1: load with randomly initialized weights (train from scratch) | |
| #tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B") | |
| #model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B") | |
| config = ViTConfig(num_hidden_layers=12, hidden_size=768) | |
| model = ViTForImageClassification(config) | |
| #print(config) | |
| feature_extractor = ViTFeatureExtractor() | |
| # or, to load one that corresponds to a checkpoint on the hub: | |
| #feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224") | |
| #the following gets called by classify_image() | |
| feature_extractor = PerceiverFeatureExtractor.from_pretrained("deepmind/vision-perceiver-conv") | |
| model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv") | |
| #google/vit-base-patch16-224, deepmind/vision-perceiver-conv | |
| image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor) | |
| def inference(input_sentence, max_length, sample_or_greedy, seed=42): | |
| if sample_or_greedy == "Sample": | |
| parameters = { | |
| "max_new_tokens": max_length, | |
| "top_p": 0.9, | |
| "do_sample": True, | |
| "seed": seed, | |
| "early_stopping": False, | |
| "length_penalty": 0.0, | |
| "eos_token_id": None, | |
| } | |
| else: | |
| parameters = { | |
| "max_new_tokens": max_length, | |
| "do_sample": False, | |
| "seed": seed, | |
| "early_stopping": False, | |
| "length_penalty": 0.0, | |
| "eos_token_id": None, | |
| } | |
| payload = {"inputs": input_sentence, "parameters": parameters,"options" : {"use_cache": False} } | |
| data = query(payload) | |
| if "error" in data: | |
| return (None, None, f"<span style='color:red'>ERROR: {data['error']} </span>") | |
| generation = data[0]["generated_text"].split(input_sentence, 1)[1] | |
| return ( | |
| before_prompt | |
| + input_sentence | |
| + prompt_to_generation | |
| + generation | |
| + after_generation, | |
| data[0]["generated_text"], | |
| "", | |
| ) | |
| def create_story(text_seed): | |
| #tokenizer = AutoTokenizer.from_pretrained("gpt2") | |
| #model = AutoModelForCausalLM.from_pretrained("gpt2") | |
| #eleutherAI gpt-3 based | |
| tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M") | |
| model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M") | |
| # set pad_token_id to eos_token_id because GPT2 does not have a EOS token | |
| model.config.pad_token_id = model.config.eos_token_id | |
| #input_prompt = "It might be possible to" | |
| input_prompt = text_seed | |
| input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids | |
| # instantiate logits processors | |
| logits_processor = LogitsProcessorList( | |
| [ | |
| MinLengthLogitsProcessor(10, eos_token_id=model.config.eos_token_id), | |
| ] | |
| ) | |
| stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=100)]) | |
| outputs = model.greedy_search( | |
| input_ids, logits_processor=logits_processor, stopping_criteria=stopping_criteria | |
| ) | |
| result_text = tokenizer.batch_decode(outputs, skip_special_tokens=True) | |
| return result_text | |
| def self_caption(image): | |
| repo_name = "ydshieh/vit-gpt2-coco-en" | |
| #test_image = "cats.jpg" | |
| test_image = image | |
| #url = 'http://images.cocodataset.org/val2017/000000039769.jpg' | |
| #test_image = Image.open(requests.get(url, stream=True).raw) | |
| #test_image.save("cats.png") | |
| feature_extractor2 = ViTFeatureExtractor.from_pretrained(repo_name) | |
| tokenizer = AutoTokenizer.from_pretrained(repo_name) | |
| model2 = VisionEncoderDecoderModel.from_pretrained(repo_name) | |
| pixel_values = feature_extractor2(test_image, return_tensors="pt").pixel_values | |
| print("Pixel Values") | |
| print(pixel_values) | |
| # autoregressively generate text (using beam search or other decoding strategy) | |
| generated_ids = model2.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True) | |
| # decode into text | |
| preds = tokenizer.batch_decode(generated_ids[0], skip_special_tokens=True) | |
| preds = [pred.strip() for pred in preds] | |
| print("Predictions") | |
| print(preds) | |
| print("The preds type is : ",type(preds)) | |
| pred_keys = ["Prediction"] | |
| pred_value = preds | |
| pred_dictionary = dict(zip(pred_keys, pred_value)) | |
| print("Pred dictionary") | |
| print(pred_dictionary) | |
| #return(pred_dictionary) | |
| preds = ' '.join(preds) | |
| story = create_story(preds) | |
| story = ' '.join(story) | |
| return story | |
| def classify_image(image): | |
| results = image_pipe(image) | |
| print("RESULTS") | |
| print(results) | |
| # convert to format Gradio expects | |
| output = {} | |
| for prediction in results: | |
| predicted_label = prediction['label'] | |
| score = prediction['score'] | |
| output[predicted_label] = score | |
| print("OUTPUT") | |
| print(output) | |
| return output | |
| image = gr.inputs.Image(type="pil") | |
| label = gr.outputs.Label(num_top_classes=5) | |
| examples = [ ["cats.jpg"], ["batter.jpg"],["drinkers.jpg"] ] | |
| title = "Generate a Story from an Image" | |
| description = "Demo for classifying images with Perceiver IO. To use it, simply upload an image and click 'submit', a story is autogenerated as well" | |
| article = "<p style='text-align: center'></p>" | |
| img_info1 = gr.Interface( | |
| fn=classify_image, | |
| inputs=image, | |
| outputs=label, | |
| ) | |
| img_info2 = gr.Interface( | |
| fn=self_caption, | |
| inputs=image, | |
| #outputs=label, | |
| outputs = [ | |
| gr.outputs.Textbox(label = 'Story') | |
| ], | |
| ) | |
| Parallel(img_info1,img_info2, inputs=image, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True) | |
| #Parallel(img_info1,img_info2, inputs=image, outputs=label, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True) | |