Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from pdf2image import convert_from_path | |
| from qwen2_inference import run_inference | |
| from util import compile_tikz_to_pdf | |
| args = {} | |
| # Sidebar Setup | |
| st.sidebar.title("Model Configuration") | |
| model_name = st.sidebar.selectbox("Model Name", ['Itsumi-st/Imgtikz_Qwen2vl', 'Qwen/Qwen2-VL-7B-Instruct']) | |
| args['inference_strat'] = st.sidebar.selectbox("Inference Strategy", ["Iterative", "Multi-candidate"], | |
| help="Choose the inference strategy for the model. Iterative generates one candidate at a time until an output compiles, while Multi-candidate generates multiple candidates in parallel.") | |
| args['max_length'] = st.sidebar.slider("Max Length", 1, 5096, 2048, | |
| help="Maximum length of the generated output. The model will generate text up to this length.") | |
| args['seed'] = st.sidebar.number_input("Seed", min_value=0, value=42, step=1) | |
| args['temperature'] = st.sidebar.slider("Temperature", 0.0, 1.0, 0.6, step=0.01, | |
| help="Temperature parameter for sampling. Higher values result in more random outputs.") | |
| args['top_p'] = st.sidebar.slider("Top P", 0.0, 1.0, 1.0, step=0.01, | |
| help="Top P sampling parameter. The model will sample from the top P percentage of the probability distribution.") | |
| args['top_k'] = st.sidebar.slider("Top K", 0, 100, 50, step=1, | |
| help="Top K sampling parameter. The model will sample from the top K tokens with the highest probabilities.") | |
| # Introduction Section | |
| st.title("Sketch2Diagram") | |
| st.write("This is a runnable demo of ImgTikZ model introduced in the Sketch2Diagram paper.") | |
| st.write("Please refer to the [original paper](https://openreview.net/pdf?id=KvaDHPhhir) for more details.") | |
| st.write("The model is trained to convert sketches into TikZ code, which can be used to generate vectorized diagrams.") | |
| # User Input Section | |
| st.subheader("Upload your sketch") | |
| input_method = st.selectbox("Input Method", ["Upload", "Camera"], | |
| help="Choose how you want to input your sketch. You can either upload an image or take a picture using your webcam.") | |
| input_file = None | |
| if input_method == "Camera": | |
| input_file = st.camera_input("Take a picture of your sketch") | |
| # todo: Implement camera input functionality here | |
| else: | |
| input_file = st.file_uploader("Upload an image of your sketch", type=["png", "jpg", "jpeg"]) | |
| st.write(args) | |
| generate_command = None | |
| # Display the uploaded image | |
| if input_file is not None: | |
| st.image(input_file, caption="Uploaded Sketch") | |
| generate_command = st.button("Generate TikZ Code") | |
| # Run model inference | |
| if generate_command: | |
| with st.spinner("Generating TikZ code..."): | |
| try: | |
| output = run_inference(input_file, model_name, args)[0] | |
| except Exception as e: | |
| st.error(f"Inference failed: {e}") | |
| st.session_state.clear() | |
| return | |
| pdf_file_path = compile_tikz_to_pdf(output) | |
| if output and pdf_file_path: | |
| st.subheader("Generated TikZ Code") | |
| st.success("TikZ code generated successfully!") | |
| st.code(output, language='latex') | |
| st.download_button( | |
| label="Download LaTeX Code", | |
| data=output, | |
| file_name="output.tex", | |
| mime="text/plain" | |
| ) | |
| st.subheader("Generated Diagram") | |
| images = convert_from_path(pdf_file_path) | |
| st.image(images[0], caption="Generated Diagram", use_container_width=True) | |
| with open(pdf_file_path, "rb") as f: | |
| st.download_button( | |
| label="Download PDF", | |
| data=f.read(), | |
| file_name="output.pdf", | |
| mime="application/pdf" | |
| ) | |
| else: | |
| st.error("Failed to generate TikZ code.") | |