Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from huggingface_hub import InferenceClient, AsyncInferenceClient | |
| from PIL import Image | |
| from pathlib import Path | |
| import os, subprocess | |
| import requests | |
| st.set_page_config(page_title='HF Inference Client Demo',layout="wide") | |
| # Cache the header of the app to prevent re-rendering on each load | |
| def display_app_header(): | |
| """Display the header of the Streamlit app.""" | |
| st.title("HF Inference Client Demo") | |
| st.subheader("Just a little demonstrator") | |
| # Display the header of the app | |
| display_app_header() | |
| # UI sidebar parameters #################################### | |
| st.sidebar.header("Loging") | |
| if hg_token :=st.sidebar.text_input('Enter your HF token', type="password"): | |
| st.sidebar.info('Logged', icon="ℹ️") | |
| else: | |
| st.sidebar.warning("enter your token") | |
| st.sidebar.header("Model") | |
| selected_model = st.sidebar.radio( | |
| "Choose a model or let the client do it", | |
| ["Not choose", "Choose"] | |
| ) | |
| if selected_model == "Choose": | |
| model = st.sidebar.text_input('Enter a model name. ex : facebook/fastspeech2-en-ljspeech') | |
| else: | |
| model = None | |
| st.sidebar.header("Task") | |
| dict_hg_tasks = { | |
| "Automatic Speech Recognition":"automatic_speech_recognition", | |
| "Text-to-Speech (choose model)":"text_to_speech", | |
| "Image Classification":"image_classification", | |
| "Image Segmentation":"image_segmentation", | |
| "Object Detection":"object_detection", | |
| "Text-to-Image":"text_to_image", | |
| "Visual Question Answering":"visual_question_answering", | |
| "Conversational":"conversational", | |
| "Feature Extraction":"feature_extraction", | |
| "Question Answering":"question_answering", | |
| "Summarization":"summarization", | |
| "Text Classification":"text_classification", | |
| "Text Generation":"text_generation", | |
| "Token Classification":"token_classification", | |
| "Translation (choose model)":"translation", | |
| } | |
| dict_hg_tasks_params = { | |
| "automatic_speech_recognition": { | |
| "input": "upload,url", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "text_to_speech": { | |
| "input": "text", | |
| "output": "audio", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "image_classification": { | |
| "input": "upload,url", | |
| "output": "image,text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "image_segmentation": { | |
| "input": "upload,url", | |
| "output": "image,text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "object_detection": { | |
| "input": "upload,url", | |
| "output": "image,text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "text_to_image": { | |
| "input": "text", | |
| "output": "image", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "visual_question_answering": { | |
| "input": "upload,url", | |
| "output": "image,text", | |
| "prompt": True, | |
| "context": False | |
| }, | |
| "image_to_image": { | |
| "input": "upload,url", | |
| "output": "image,text", | |
| "prompt": True, | |
| "context": False | |
| }, | |
| "feature_extraction": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "conversational": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "question_answering": { | |
| "input": None, | |
| "output": "text", | |
| "prompt": True, | |
| "context": True | |
| }, | |
| "text_classification": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "token_classification": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "text_generation": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "text_classification": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "translation": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| "summarization": { | |
| "input": "text", | |
| "output": "text", | |
| "prompt": False, | |
| "context": False | |
| }, | |
| } | |
| selected_task = st.sidebar.radio( | |
| "Choose the task you want to do", # see https://huggingface.co/docs/huggingface_hub/guides/inference" | |
| dict_hg_tasks.keys() | |
| ) | |
| st.write(f"The current selected task is : {dict_hg_tasks[selected_task]}") | |
| with st.sidebar.expander("tasks documentation"): | |
| st.write("https://huggingface.co/docs/huggingface_hub/package_reference/inference_client") | |
| # functions ######################################## | |
| cwd = os.getcwd() | |
| def get_input(upload,url,text): | |
| if upload is not None: | |
| return upload | |
| else: | |
| if url: | |
| print(url) | |
| return url | |
| elif text: | |
| return text | |
| return None # Default return if neither upload nor url is provided | |
| def display_inputs(task): | |
| if dict_hg_tasks_params[task]["input"] == "upload,url": | |
| return st.file_uploader("Choose a file"),st.text_input("or enter a file url"),"" | |
| elif dict_hg_tasks_params[task]["input"] == "text": | |
| return None,"",st.text_input("Enter a text") | |
| else: | |
| return None,"","" | |
| def display_prompt(task): | |
| if dict_hg_tasks_params[task]["prompt"] is True: | |
| return st.text_input("Enter a question") | |
| return None | |
| def display_context(task): | |
| if dict_hg_tasks_params[task]["context"] is True: | |
| return st.text_area("Enter a context") | |
| return None | |
| # UI main client #################################### | |
| if selected_task : | |
| response = None | |
| task = dict_hg_tasks[selected_task] | |
| if model: | |
| client = InferenceClient(model=model) | |
| else: | |
| client = InferenceClient() | |
| uploaded_input,url_input,text_input = display_inputs(task) | |
| prompt_input = display_prompt(task) | |
| context_input = display_context(task) | |
| if get_input(uploaded_input,url_input,text_input): | |
| input = get_input(uploaded_input,url_input,text_input) | |
| if prompt_input is not None: | |
| response = getattr(client, task)(input,question=prompt_input) | |
| else: | |
| response = getattr(client, task)(input) | |
| elif prompt_input: | |
| if context_input is not None: | |
| response = getattr(client, task)(question=prompt_input,context=context_input) | |
| else: | |
| response = getattr(client, task)(input,prompt=prompt_input) | |
| if response is not None: | |
| col1,col2 = st.columns(2) | |
| with col1: | |
| if "text" in dict_hg_tasks_params[task]["output"]: | |
| st.write(response) | |
| elif "audio" in dict_hg_tasks_params[task]["output"]: | |
| Path(os.path.join(cwd,"audio.flac")).write_bytes(response) | |
| st.audio(os.path.join(cwd,"audio.flac")) | |
| with col2: | |
| if dict_hg_tasks_params[task]["output"] == "image,text": | |
| if not(isinstance(input, str)): | |
| image = Image.open(input) | |
| else: | |
| image = Image.open(requests.get(input, stream=True).raw) | |
| st.image(image) | |
| elif dict_hg_tasks_params[task]["output"] == "image": | |
| response.save(os.path.join(cwd,"generated_image.png")) | |
| image = Image.open(os.path.join(cwd,"generated_image.png")) | |
| st.image(image) |