Update app.py

app.py

@@ -2,30 +2,75 @@ import streamlit as st
 from diffusers import StableDiffusionPipeline
 import torch
 from PIL import Image
+import tempfile
+import whisper
+import os
+from realesrgan import RealESRGAN
 
+# Load Whisper for voice-to-text
 @st.cache_resource
-def load_pipeline():
-    model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float32,
-        use_safetensors=True
-    )
-    return pipe.to("cpu")  # if you're using CPU on Hugging Face Spaces
+def load_whisper():
+    return whisper.load_model("base")
 
-pipe = load_pipeline()
+# Load Stable Diffusion models
+@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
+
+# Load Real-ESRGAN for enhancement
+@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.title("💖 Love Text to Image Generator")
+
+st.markdown("Write or speak a romantic poem and turn it into beautiful art ✨")
 
-st.markdown("Write a poem or love message, and see it transformed into an image!")
+style = st.selectbox("Choose art style", ["Realistic", "Anime", "Ghibli"])
 
-user_input = st.text_area("Enter your love poem, song, or message here...", height=150)
+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 user_input.strip() == "":
-        st.warning("Please write something romantic to generate an image!")
+    if prompt.strip() == "":
+        st.warning("Please enter or upload a love message.")
     else:
         with st.spinner("Generating image..."):
-            image = pipe(user_input).images[0]
-            st.image(image, caption="Generated from your text 💖", use_column_width=True)
+            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)