Spaces:
Paused
Paused
| import streamlit as st | |
| import base64 | |
| from ml import MLModel | |
| from naive import NaiveModel | |
| import torch | |
| st.set_page_config(page_title="Drawing with LLM", page_icon="🎨", layout="wide") | |
| def load_ml_model(): | |
| return MLModel(device="cuda" if st.session_state.get("use_gpu", True) else "cpu") | |
| def load_naive_model(): | |
| return NaiveModel(device="cuda" if st.session_state.get("use_gpu", True) else "cpu") | |
| def render_svg(svg_content): | |
| b64 = base64.b64encode(svg_content.encode("utf-8")).decode("utf-8") | |
| return f'<img src="data:image/svg+xml;base64,{b64}" width="100%" height="auto"/>' | |
| def clear_gpu_memory(): | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| torch.cuda.ipc_collect() | |
| st.title("Drawing with LLM 🎨") | |
| # Initialize session state for model type if not already set | |
| if "current_model_type" not in st.session_state: | |
| st.session_state["current_model_type"] = None | |
| with st.sidebar: | |
| st.header("Settings") | |
| previous_model_type = st.session_state.get("current_model_type") | |
| model_type = st.selectbox("Model Type", ["ML Model (vtracer)", "Naive Model (phi-4)"]) | |
| # Check if model type has changed | |
| if previous_model_type is not None and previous_model_type != model_type: | |
| st.cache_resource.clear() | |
| clear_gpu_memory() | |
| st.success(f"Cleared VRAM after switching from {previous_model_type} to {model_type}") | |
| # Update current model type in session state | |
| st.session_state["current_model_type"] = model_type | |
| use_gpu = st.checkbox("Use GPU", value=True) | |
| st.session_state["use_gpu"] = use_gpu | |
| if model_type == "ML Model (vtracer)": | |
| st.subheader("ML Model Settings") | |
| simplify = st.checkbox("Simplify SVG", value=True) | |
| color_precision = st.slider("Color Precision", 1, 10, 6) | |
| filter_speckle = st.slider("Filter Speckle", 0, 10, 4) | |
| path_precision = st.slider("Path Precision", 1, 10, 8) | |
| elif model_type == "Naive Model (phi-4)": | |
| st.subheader("Naive Model Settings") | |
| max_new_tokens = st.slider("Max New Tokens", 256, 1024, 512) | |
| prompt = st.text_area("Enter your description", "A cat sitting on a windowsill at sunset") | |
| if st.button("Generate SVG"): | |
| with st.spinner("Generating SVG..."): | |
| if model_type == "ML Model (vtracer)": | |
| model = load_ml_model() | |
| svg_content = model.predict( | |
| prompt, | |
| simplify=simplify, | |
| color_precision=color_precision, | |
| filter_speckle=filter_speckle, | |
| path_precision=path_precision | |
| ) | |
| else: # Naive Model | |
| model = load_naive_model() | |
| svg_content = model.predict(prompt, max_new_tokens=max_new_tokens) | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.subheader("Generated SVG") | |
| st.markdown(render_svg(svg_content), unsafe_allow_html=True) | |
| with col2: | |
| st.subheader("SVG Code") | |
| st.code(svg_content, language="xml") | |
| # Download button for SVG | |
| st.download_button( | |
| label="Download SVG", | |
| data=svg_content, | |
| file_name="generated_svg.svg", | |
| mime="image/svg+xml" | |
| ) | |
| st.markdown("---") | |
| st.markdown("This app uses Stable Diffusion to generate images from text and converts them to SVG.") | |