0ddly/Memealot

1import streamlit as st2import torch3from diffusers import StableDiffusionPipeline4from PIL import Image5import os6 7def generate_image(prompt):8  """9  Generate an image based on the given prompt.10  11  Args:12  prompt (str): The text prompt to generate an image from.13  14  Returns:15  PIL.Image: The generated image.16  """17  try:18    device = "cuda" if torch.cuda.is_available() else "cpu"19    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True, device=device)20    torch.cuda.empty_cache()  # Free up GPU memory21    pipe.disable_progress_bar()  # Disable progress bar22 23    image = pipe(prompt).images[0]24    return image25  except ConnectionError as e:26    st.error(f"Error loading model: {e}")27    return None28  except TimeoutError as e:29    st.error(f"Timeout error: {e}")30    return None31  except Exception as e:32    st.error(f"Error generating image: {e}")33    return None34 35def save_image(image, filename):36  """37  Save the generated image to a file.38  39  Args:40  image (PIL.Image): The generated image.41  filename (str): The filename to save the image as.42  """43  if image is None:44    print("No image to save.")45    return46  try:47    if os.path.exists(filename):48      st.error("File already exists. Please enter a different filename.")49      return50    image.save(filename, format="JPEG")  # Specify JPEG format51    st.success(f"Image saved as '{filename}'")52  except Exception as e:53    st.error54