Update app.py

app.py

@@ -5,30 +5,31 @@ from diffusers import DiffusionPipeline
 from transformers import pipeline
 
 # Load text generation pipeline
-text_pipe = pipeline('text-generation', model='daspartho/prompt-extend')
+@st.cache_resource
+def load_text_pipeline():
+    return pipeline('text-generation', model='daspartho/prompt-extend')
+
+text_pipe = load_text_pipeline()
 
 def extend_prompt(prompt):
     return text_pipe(prompt + ',', num_return_sequences=1, max_new_tokens=50)[0]["generated_text"]
 
 @st.cache_resource
-def load_pipeline(use_cuda):
-    device = "cuda" if use_cuda and torch.cuda.is_available() else "cpu"
+def load_pipeline():
     pipe = DiffusionPipeline.from_pretrained(
         "stabilityai/sdxl-turbo",
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        variant="fp16" if device == "cuda" else None,
+        torch_dtype=torch.float32,
+        variant=None,
         use_safetensors=True
     )
-    if device == "cuda":
-        pipe.enable_xformers_memory_efficient_attention()
-    pipe.to(device)
+    pipe.to("cpu")
     return pipe
 
-def generate_image(prompt, use_details, use_cuda):
-    pipe = load_pipeline(use_cuda)
+def generate_image(prompt, use_details):
+    pipe = load_pipeline()
     generator = torch.manual_seed(np.random.randint(0, 2**32))
     extended_prompt = extend_prompt(prompt) if use_details else prompt
-    image = pipe(prompt=extended_prompt, generator=generator, num_inference_steps=25, guidance_scale=7.5).images[0]
+    image = pipe(prompt=extended_prompt, generator=generator, num_inference_steps=15, guidance_scale=7.5).images[0]
     return image, extended_prompt
 
 # Add the custom CSS file
@@ -41,11 +42,10 @@ st.markdown("<div class='subheader'>Create an anime-style GitHub profile picture
 input_text = st.text_area("Describe your GitHub profile picture:", "Create an anime-style GitHub profile picture for a boy")
 
 details_checkbox = st.checkbox("Generate Details?", value=True)
-cuda_checkbox = st.checkbox("Use CUDA?", value=False)
 
 if st.button("Generate Image"):
     with st.spinner('Generating image...'):
-        image, extended_prompt = generate_image(input_text, details_checkbox, cuda_checkbox)
+        image, extended_prompt = generate_image(input_text, details_checkbox)
         st.image(image, caption="Generated Image")
         st.text(f"Extended Prompt: {extended_prompt}")
     st.balloons()