Spaces:
Paused
Paused
File size: 3,468 Bytes
da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 8bf4ef4 da1ac10 |
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 86 87 88 89 90 91 92 93 94 95 96 |
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")
@st.cache_resource
def load_ml_model():
return MLModel(device="cuda" if st.session_state.get("use_gpu", True) else "cpu")
@st.cache_resource
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.")
|