Update app.py

app.py

@@ -1,84 +1,92 @@
+# app.py
 import streamlit as st
-import os
-import cv2
-import numpy as np
-from PIL import Image
-from rembg import remove
-from realesrgan import RealESRGAN
-from moviepy.editor import ImageSequenceClip
-
-# Upload section
-st.title("🖼️ AI Image & Video Effects Studio")
-st.markdown("Upload multiple images, apply cool effects, and export a video!")
-
-uploaded_files = st.file_uploader("Upload images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
-
-# Select effect
-effect = st.selectbox("Choose an effect to apply", ["None", "Cartoon", "Edge Detection", "Remove Background", "Face Enhancement"])
-
-output_dir = "outputs"
-os.makedirs(output_dir, exist_ok=True)
-
-processed_images = []
-
-def cartoonize(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    gray = cv2.medianBlur(gray, 5)
-    edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
-                                  cv2.THRESH_BINARY, 9, 9)
-    color = cv2.bilateralFilter(image, 9, 300, 300)
-    cartoon = cv2.bitwise_and(color, color, mask=edges)
-    return cartoon
-
-def edge_detect(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    edges = cv2.Canny(gray, 100, 200)
-    return cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
-
-def enhance_face(image_pil):
-    from realesrgan import RealESRGAN
-    model = RealESRGAN("cpu", scale=2)
-    model.load_weights("RealESRGAN_x2.pth")
-    return model.predict(image_pil)
-
-def save_video(images, filename="output_video.mp4"):
-    clip = ImageSequenceClip(images, fps=1)
-    clip.write_videofile(filename, codec="libx264")
-    return filename
-
-if uploaded_files:
-    for file in uploaded_files:
-        img = Image.open(file).convert("RGB")
-        np_img = np.array(img)
-        np_img_bgr = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
-
-        if effect == "Cartoon":
-            result = cartoonize(np_img_bgr)
-        elif effect == "Edge Detection":
-            result = edge_detect(np_img_bgr)
-        elif effect == "Remove Background":
-            result = remove(np_img)
-        elif effect == "Face Enhancement":
-            result = enhance_face(img)
-        else:
-            result = np_img_bgr
-
-        result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) if isinstance(result, np.ndarray) else result
-        final_img = Image.fromarray(result_rgb)
-
-        st.image(final_img, caption="Processed Image", use_column_width=True)
-
-        img_path = os.path.join(output_dir, f"{file.name}")
-        final_img.save(img_path)
-        processed_images.append(img_path)
-
-    # Convert to video
-    if st.button("🎞️ Convert to Video Slideshow"):
-        video_path = save_video(processed_images)
-        st.video(video_path)
-
-    # Download
-    for img_path in processed_images:
-        with open(img_path, "rb") as f:
-            st.download_button("Download Image", f, file_name=os.path.basename(img_path))
-
+from diffusers import StableDiffusionPipeline
+import torch
+
+# Cache the model loading to avoid reloading on every interaction
+@st.cache_resource
+def load_model():
+    model_id = "CompVis/stable-diffusion-v1-4"
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32,
+        use_safetensors=True,
+        safety_checker=None
+    )
+    pipe = pipe.to("cpu")
+    pipe.enable_attention_slicing()  # Reduce memory usage
+    return pipe
+
+def main():
+    st.set_page_config(page_title="Poetry to Image", page_icon="🎨")
+    
+    st.title("✨ Romantic Text to Image Converter")
+    st.markdown("Turn your poetic words into beautiful images (CPU version)")
+    
+    # Initialize session state
+    if 'generated_image' not in st.session_state:
+        st.session_state.generated_image = None
+    if 'error' not in st.session_state:
+        st.session_state.error = None
+    
+    # Sidebar with settings
+    with st.sidebar:
+        st.header("Settings")
+        num_steps = st.slider("Inference Steps", 10, 50, 25)
+        guidance_scale = st.slider("Guidance Scale", 3.0, 15.0, 7.5)
+    
+    # Main interface
+    prompt = st.text_area(
+        "Enter your romantic text or poetry:",
+        height=150,
+        placeholder="Example: 'Your eyes sparkle like stars in the night sky...'"
+    )
+    
+    col1, col2 = st.columns([1, 3])
+    with col1:
+        if st.button("Generate Image", type="primary"):
+            st.session_state.generated_image = None
+            st.session_state.error = None
+            
+            if not prompt.strip():
+                st.session_state.error = "Please enter some text first!"
+            else:
+                try:
+                    with st.spinner("Creating your artwork... (this may take a few minutes)"):
+                        pipe = load_model()
+                        image = pipe(
+                            prompt,
+                            num_inference_steps=num_steps,
+                            guidance_scale=guidance_scale
+                        ).images[0]
+                        
+                        st.session_state.generated_image = image
+                except Exception as e:
+                    st.session_state.error = f"Error: {str(e)}"
+    
+    # Display results
+    if st.session_state.error:
+        st.error(st.session_state.error)
+    
+    if st.session_state.generated_image:
+        st.subheader("Generated Image")
+        st.image(
+            st.session_state.generated_image,
+            use_column_width=True,
+            caption="Your generated artwork"
+        )
+        
+        # Add download button
+        buf = io.BytesIO()
+        st.session_state.generated_image.save(buf, format="PNG")
+        byte_im = buf.getvalue()
+        
+        st.download_button(
+            label="Download Image",
+            data=byte_im,
+            file_name="poetry_image.png",
+            mime="image/png"
+        )
+
+if __name__ == "__main__":
+    main()