import streamlit as st
from llama_cpp import Llama
from PIL import Image
import numpy as np
from transformers import AutoModel
model = AutoModel.from_pretrained("liminerity/bitmap-mistral-M7-slerp-alpaca-70m-gguf")

# Initialize the model (only once)
@st.cache_resource
def load_model():
    return Llama(
        model_path=model,
        n_ctx=2048,
        n_threads=4
    )

st.title("BitMap Generator")

# Input text prompt
prompt = st.text_input("Enter your prompt for the bitmap:", "A simple circle")

if st.button("Generate Bitmap"):
    if prompt:
        # Generate the bitmap description
        llm = load_model()
        response = llm(
            f"Generate a 64x64 bitmap array using 0s and 1s that represents {prompt}. "
            "Only output the array, no other text.",
            max_tokens=2048,
            temperature=0.7
        )
        
        # Convert the response to a bitmap
        try:
            # Parse the response to get just the array
            array_text = response['choices'][0]['text'].strip()
            # Convert string to numpy array
            bitmap = np.array([list(map(int, row.strip('[] ').split()))
                             for row in array_text.split('\n') if row.strip()])
            
            # Scale up the bitmap for better visibility
            scaled_bitmap = np.kron(bitmap, np.ones((10, 10)))
            
            # Create an image from the array
            img = Image.fromarray(np.uint8(scaled_bitmap * 255))
            
            # Display the image
            st.image(img, caption=f"Generated bitmap for: {prompt}")
        except Exception as e:
            st.error(f"Error creating bitmap: {str(e)}")