File size: 3,959 Bytes
30f49a1
8557bbe
30f49a1
979c542
8557bbe
979c542
 
30f49a1
 
 
979c542
 
8557bbe
 
 
979c542
8557bbe
 
 
 
 
 
30f49a1
 
 
 
 
 
 
979c542
30f49a1
 
 
 
 
 
 
 
 
8557bbe
30f49a1
 
8557bbe
30f49a1
 
 
 
 
 
979c542
30f49a1
 
9e90829
 
 
 
 
0c3ac3e
8557bbe
 
e497738
8557bbe
 
 
 
 
 
 
 
 
 
e497738
 
722d92d
8557bbe
 
 
e497738
8557bbe
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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.")