|
|
import streamlit as st |
|
|
from diffusers import StableDiffusionPipeline |
|
|
import torch |
|
|
from PIL import Image |
|
|
import tempfile |
|
|
import whisper |
|
|
import os |
|
|
from realesrgan import RealESRGAN |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
|
|
|
@st.cache_resource |
|
|
def load_whisper(): |
|
|
return whisper.load_model("base") |
|
|
|
|
|
|
|
|
@st.cache_resource |
|
|
def load_pipelines(): |
|
|
pipelines = { |
|
|
"Realistic": StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32), |
|
|
"Anime": StableDiffusionPipeline.from_pretrained("andite/anything-v4.0", torch_dtype=torch.float32), |
|
|
"Ghibli": StableDiffusionPipeline.from_pretrained("nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float32) |
|
|
} |
|
|
for key in pipelines: |
|
|
pipelines[key] = pipelines[key].to("cpu") |
|
|
return pipelines |
|
|
|
|
|
|
|
|
@st.cache_resource |
|
|
def load_enhancer(): |
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet |
|
|
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, |
|
|
num_block=23, num_grow_ch=32, scale=4) |
|
|
return RealESRGAN(device=torch.device("cpu"), scale=4).load_weights(model) |
|
|
|
|
|
whisper_model = load_whisper() |
|
|
pipes = load_pipelines() |
|
|
enhancer = load_enhancer() |
|
|
|
|
|
st.set_page_config(page_title="Love Text to Image", layout="centered") |
|
|
st.title("💖 Love Text to Image Generator") |
|
|
|
|
|
st.markdown("Write or speak a romantic poem and turn it into beautiful art ✨") |
|
|
|
|
|
style = st.selectbox("Choose art style", ["Realistic", "Anime", "Ghibli"]) |
|
|
|
|
|
tab1, tab2 = st.tabs(["📝 Text Input", "🎤 Voice Input"]) |
|
|
|
|
|
prompt = "" |
|
|
|
|
|
with tab1: |
|
|
prompt = st.text_area("Enter your romantic poem or message") |
|
|
|
|
|
with tab2: |
|
|
audio_file = st.file_uploader("Upload your audio file (mp3, wav)", type=["mp3", "wav"]) |
|
|
if audio_file is not None: |
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp: |
|
|
tmp.write(audio_file.read()) |
|
|
tmp_path = tmp.name |
|
|
with st.spinner("Transcribing..."): |
|
|
result = whisper_model.transcribe(tmp_path) |
|
|
prompt = result["text"] |
|
|
st.success("Transcription Complete!") |
|
|
st.markdown(f"**Transcribed Text:** {prompt}") |
|
|
os.remove(tmp_path) |
|
|
|
|
|
if st.button("Generate Image"): |
|
|
if prompt.strip() == "": |
|
|
st.warning("Please enter or upload a love message.") |
|
|
else: |
|
|
with st.spinner("Generating image..."): |
|
|
image = pipes[style](prompt).images[0] |
|
|
st.image(image, caption="Original Generated Image", use_column_width=True) |
|
|
|
|
|
with st.spinner("Enhancing image..."): |
|
|
enhanced = enhancer.predict(image) |
|
|
st.image(enhanced, caption="Enhanced Image", use_column_width=True) |
|
|
|
