Reinstate model caching

- uncomment @cache_resource to reactivate model caching
- Added subheaders for generated output

main.py

@@ -1,5 +1,4 @@
 import streamlit as st
-from transformers import pipeline
 
 st.logo("NLP_Group_logo.svg", size="large")
 main_page = st.Page("main_page.py", title="Main Page", icon="🏠")
@@ -8,9 +7,3 @@ sketch2diagram_page = st.Page("sketch2diagram.py", title="Sketch2Diagram", icon=
 pg = st.navigation([main_page, sketch2diagram_page])
 
 pg.run()
-
-@st.cache_resource
-def get_model():
-    # Load the model here
-    model = pipeline("image-to-text", model="itsumi-st/imgtikz_qwen2vl")
-    return model

qwen2_inference.py

@@ -20,8 +20,7 @@ def print_gpu_memory(label, memory_allocated, memory_reserved):
 
 
 # Inference steps taken from https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct
-
-# @st.cache_resource
+@st.cache_resource
 def get_model(model_path):
     try:
         with st.spinner(f"Loading model {model_path}"):

sketch2diagram.py

@@ -53,6 +53,7 @@ if generate_command:
         output = run_inference(input_file, model_name, args)[0]
         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')
 
@@ -63,16 +64,16 @@ if generate_command:
                 mime="text/plain"
             )
 
-            # st.image(pdf_file_path, caption="Generated Diagram", use_column_width=True)
+            st.subheader("Generated Diagram")
+            images = convert_from_path(pdf_file_path)
+            st.image(images[0], caption="Generated Diagram", use_column_width=True)
             with open(pdf_file_path, "rb") as f:
                 st.download_button(
                     label="Download PDF",
-                    data=f.read(),  # ✅ this is the binary content
+                    data=f.read(),
                     file_name="output.pdf",
                     mime="application/pdf"
                 )
 
-            images = convert_from_path(pdf_file_path)
-            st.image(images[0], caption="Generated Diagram", use_column_width=True)
         else:
             st.error("Failed to generate TikZ code.")