Upload app.py

app.py

@@ -3,23 +3,22 @@ import torch
 from diffusers import StableDiffusionXLPipeline
 from PIL import Image
 
-# Load the model
+# Load the model (CPU version)
 @st.cache_resource()
 def load_model():
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "cagliostrolab/animagine-xl-4.0",
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float32,  # Change to float32 for CPU
         use_safetensors=True,
         custom_pipeline="lpw_stable_diffusion_xl",
         add_watermarker=False
     )
-    #pipe.to("cuda")
-    return pipe
+    return pipe  # Do NOT move to CUDA
 
 pipe = load_model()
 
 # Streamlit UI
-st.title("Text-to-Image Generator")
+st.title("Text-to-Image Generator (CPU)")
 
 # User input
 prompt = st.text_area("Enter your prompt:", "1girl, arima kana, oshi no ko, looking at viewer, smile")
@@ -29,20 +28,23 @@ guidance_scale = st.slider("Guidance Scale", 1.0, 10.0, 6.0, step=0.5)
 inference_steps = st.slider("Inference Steps", 10, 50, 25, step=5)
 
 if st.button("Generate Image"):
-    with st.spinner("Generating image..."):
-        image = pipe(
-            prompt,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=inference_steps
-        ).images[0]
-        
-        # Save and display image
+    try:
+        with st.spinner("Generating image..."):
+            image = pipe(
+                prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=inference_steps
+            ).images[0]
+
+        st.image(image, caption="Generated Image", use_column_width=True)
+
         image_path = "generated_image.png"
         image.save(image_path)
-        st.image(image, caption="Generated Image", use_column_width=True)
-        
-        # Provide download option
+
         with open(image_path, "rb") as file:
             st.download_button(label="Download Image", data=file, file_name="generated_image.png", mime="image/png")
+
+    except Exception as e:
+        st.error(f"Error generating image: {e}")