Update app.py

app.py

@@ -4,23 +4,17 @@ from diffusers import AutoPipelineForImage2Image
 from diffusers.utils import load_image
 from PIL import Image
 
-# Streamlit UI
-st.title("Stable Diffusion Image-to-Image Generation")
+st.title("Stable Diffusion Image-to-Image")
 
-# Load the pipeline
-@st.cache_resource()
-def load_pipeline():
-    pipeline = AutoPipelineForImage2Image.from_pretrained(
-        "stabilityai/stable-diffusion-2-1",  # Updated model for better quality
-        torch_dtype=torch.float16,
-        variant="fp16",
-        use_safetensors=True
-    )
-    pipeline.enable_model_cpu_offload()
-    pipeline.enable_xformers_memory_efficient_attention()
-    return pipeline
-
-pipeline = load_pipeline()
+# Load the pipeline (no caching)
+pipeline = AutoPipelineForImage2Image.from_pretrained(
+    "stabilityai/stable-diffusion-2-1",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True
+)
+pipeline.enable_model_cpu_offload()
+pipeline.enable_xformers_memory_efficient_attention()
 
 # Upload an image
 uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
@@ -29,11 +23,9 @@ if uploaded_file:
     init_image = Image.open(uploaded_file).convert("RGB")
     st.image(init_image, caption="Uploaded Image", use_column_width=True)
 
-    # Prompt input
     prompt = st.text_input("Enter a prompt", "A futuristic city at sunset")
 
-    # Generate button
     if st.button("Generate Image"):
         with st.spinner("Generating..."):
             image = pipeline(prompt, image=init_image).images[0]
-            st.image(image, caption="Generated Image", use_column_width=True) 
+            st.image(image, caption="Generated Image", use_column_width=True)