Update app.py

app.py

@@ -1,132 +1,86 @@
 # app.py
- #Add before any other imports
-from huggingface_hub import cached_download, hf_hub_download
-import warnings
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-
-import streamlit as st
-from diffusers import StableDiffusionPipeline
-import torch
-
 import streamlit as st
 import torch
+import io
+from PIL import Image
 from diffusers import StableDiffusionPipeline
 
-
-# Model configuration
+# Configuration
 MODEL_ID = "CompVis/stable-diffusion-v1-4"
+DEFAULT_STEPS = 25
+DEFAULT_GUIDANCE = 7.5
 
-@st.cache_resource(show_spinner=False)
-def load_model():
-    return StableDiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.float32,
-        use_auth_token=False,
-        revision="fp16",
-        use_safetensors=False,  # Force use of .bin files
-        custom_pipeline="stable_diffusion",
-        safety_checker=None
-    ).to("cpu")
-
-# Rest of your Streamlit code remains the same...
-# Add this at the top of your app.py
-import warnings
-warnings.filterwarnings("ignore", message="cached_download is deprecated")
-
-# Modify your model loading to use legacy features
 @st.cache_resource(show_spinner=False)
 def load_model():
     pipe = StableDiffusionPipeline.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.float32,
-        use_safetensors=False,  # Force using .bin files
-        use_auth_token=False,
-        safety_checker=None
-    )
-    # Rest of your loading code...
-# 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
+        use_safetensors=False,  # Force .bin files for compatibility
+        safety_checker=None,
+        requires_safety_checker=False
     )
     pipe = pipe.to("cpu")
-    pipe.enable_attention_slicing()  # Reduce memory usage
+    pipe.enable_attention_slicing()
     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
+    # Session state initialization
     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
+
+    # Sidebar controls
     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)
-    
+        num_steps = st.slider("Inference Steps", 10, 50, DEFAULT_STEPS)
+        guidance_scale = st.slider("Guidance Scale", 3.0, 15.0, DEFAULT_GUIDANCE)
+
     # Main interface
     prompt = st.text_area(
-        "Enter your romantic text or poetry:",
+        "Enter your romantic text:",
         height=150,
-        placeholder="Example: 'Your eyes sparkle like stars in the night sky...'"
+        placeholder="Example: 'Your eyes sparkle like stars...'"
     )
-    
-    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)}"
-    
+
+    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!"
+        else:
+            try:
+                with st.spinner("Creating artwork (3-8 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"
-        )
+        st.image(st.session_state.generated_image, use_column_width=True)
         
-        # Add download button
+        # Download functionality
         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"
+            "Download Image",
+            buf.getvalue(),
+            "poetry_image.png",
+            "image/png"
         )
 
 if __name__ == "__main__":