Update sketch2diagram.py

sketch2diagram.py

@@ -47,7 +47,6 @@ 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:
@@ -55,30 +54,30 @@ if generate_command:
         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"
-            )
+            # return   # <-- REMOVE THIS LINE!
+        else:
+            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.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"
+                    label="Download LaTeX Code",
+                    data=output,
+                    file_name="output.tex",
+                    mime="text/plain"
                 )
 
-        else:
-            st.error("Failed to generate TikZ code.")
+                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"
+                    )
+