Spaces:
Configuration error
Configuration error
File size: 5,806 Bytes
479cae1 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # import streamlit as st
# import requests
# st.title("FASTAPI Multi-Modal ChatBot")
# # Mode selector
# mode = st.sidebar.selectbox(
# "Choose Generation Mode",
# ["Text", "Audio", "Image","Video"]
# )
# if "messages" not in st.session_state:
# st.session_state.messages = []
# # Display previous chat history
# for message in st.session_state.messages:
# with st.chat_message(message["role"]):
# content = message["content"]
# if isinstance(content, bytes) and message.get("type") == "audio":
# st.audio(content, format="audio/mp3")
# elif message.get("type") == "image":
# st.image(content)
# else:
# st.markdown(content)
# # One chat input for all modes
# if prompt := st.chat_input(f"Enter your {mode.lower()} prompt here..."):
# st.session_state.messages.append({"role": "user", "content": prompt})
# with st.chat_message("user"):
# st.markdown(prompt)
# if mode == "Text":
# response = requests.get("http://localhost:8000/generate/text", params={"prompt": prompt})
# response.raise_for_status()
# reply = response.text
# st.session_state.messages.append({"role": "assistant", "content": reply})
# with st.chat_message("assistant"):
# st.markdown(reply)
# elif mode == "Audio":
# response = requests.get("http://localhost:8000/generate/audio", params={"prompt": prompt})
# response.raise_for_status()
# audio_bytes = response.content
# st.session_state.messages.append({"role": "assistant", "content": audio_bytes, "type": "audio"})
# with st.chat_message("assistant"):
# st.text("Here is your generated audio:")
# st.audio(audio_bytes, format="audio/mp3")
# elif mode == "Image":
# response = requests.get("http://localhost:8000/generate/image", params={"prompt": prompt})
# response.raise_for_status()
# image_bytes = response.content
# st.session_state.messages.append({"role": "assistant", "content": image_bytes, "type": "image"})
# with st.chat_message("assistant"):
# st.text("Here is your generated image:")
# st.image(image_bytes)
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from diffusers import DiffusionPipeline
st.set_page_config(page_title="TinyLLM Playground", layout="wide")
st.title("🎨 TinyLLM Multi-Modal Playground")
st.write("Text Generation | Image Generation | All running locally!")
# ============ MODEL LOADING (cached, loads once) ============
@st.cache_resource
def load_text_model():
"""Load TinyLlama for text generation"""
try:
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = AutoModelForCausalLM.from_pretrained(
"TinyLlama/TinyLlama-1.1B-Chat-v1.0",
torch_dtype=torch.float32,
device_map="cpu"
)
return tokenizer, model
except Exception as e:
st.error(f"Error loading text model: {e}")
return None, None
@st.cache_resource
def load_image_model():
"""Load Stable Diffusion for image generation"""
try:
pipe = DiffusionPipeline.from_pretrained(
"segmind/tiny-sd",
torch_dtype=torch.float32
)
pipe = pipe.to("cpu")
return pipe
except Exception as e:
st.error(f"Error loading image model: {e}")
return None
# ============ TEXT GENERATION ============
st.header("📝 Text Generation")
text_prompt = st.text_input("Enter a prompt:", placeholder="Tell me a joke about...")
if st.button("Generate Text", key="text_btn"):
if text_prompt:
with st.spinner("Generating text..."):
tokenizer, model = load_text_model()
if tokenizer and model:
try:
inputs = tokenizer(text_prompt, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
inputs["input_ids"],
max_length=150,
temperature=0.7,
top_p=0.9,
do_sample=True
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
st.success("✅ Generated!")
st.write(response)
except Exception as e:
st.error(f"Generation error: {e}")
else:
st.error("Model failed to load")
# ============ IMAGE GENERATION ============
st.header("🖼️ Image Generation")
image_prompt = st.text_input("Enter an image prompt:", placeholder="A sunset over mountains...")
if st.button("Generate Image", key="image_btn"):
if image_prompt:
with st.spinner("Generating image (this may take a minute)..."):
pipe = load_image_model()
if pipe:
try:
image = pipe(image_prompt, num_inference_steps=20).images[0]
st.success("✅ Generated!")
st.image(image, caption=f"Generated: {image_prompt}")
except Exception as e:
st.error(f"Image generation error: {e}")
else:
st.error("Image model failed to load")
st.divider()
st.info("💡 **Note:** First run downloads model weights (~2-5 GB). Subsequent runs are instant!")
st.divider()
st.info("💡 **Note:** First run downloads model weights (~2-5 GB). Subsequent runs are instant!")
|