Home » Tutorials » How to Build an Image Upscaler Using Stable Diffusion in Python

How to Build an Image Upscaler Using Stable Diffusion in Python

By /

Image upscaling can transform low-resolution visuals into high-quality images, making it a crucial tool for anyone working with photos or graphics. By leveraging the power of an image upscaler using Python, you can enhance image clarity and detail effortlessly.

Today, you’ll learn how to build an image upscaler using Python, the Stable Diffusion model, and Tkinter. With Stable Diffusion handling the upscaling process, Tkinter will allow us to create a user-friendly GUI where you can easily select an image, upscale it, and save the enhanced version. Let’s dive in and unleash the power of Python!

Table of Contents

Getting Started

Make sure you have the necessary libraries installed. You can easily install them by running the following commands in your terminal or command prompt:

$ pip install Pillow
$ pip install tk
$ pip install diffusers
$ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118  # for GPU
$ pip install torch torchvision torchaudio  # for CPU
$ pip install huggingface-hub

First things first, just like any programmer, we need to gather the tools for our project. Here’s what we’ll need:

  • Tkinter: This library will create the main window for our app, handle buttons, and manage the interface elements like message boxes and options for selecting images and saving them.
  • PIL (Pillow): It will take care of loading, resizing, and displaying the images.
  • Diffusers and Torch: These are the main components. Hugging Face’s Stable Diffusion model (from Diffusers) will upscale the chosen image, and Torch will help run the model efficiently.
  • Threading: This ensures that our program can multitask smoothly without freezing the main window.
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from tkinter import Label, Button
from PIL import Image, ImageTk
from diffusers import StableDiffusionUpscalePipeline
import torch
import threading

Setting Up the Main Window and Scrollable Interface

Creating the Main Window

With our tools ready, it’s time to set the stage. We’ll start by setting up the interface that allows users to interact with the program. To do this, we use tk to create the main window, which will hold all the widgets, and we give it a title.

# Initialize the main window
root = tk.Tk()
root.title("Image Upscaler - The Pycodes")
root.geometry("600x400")

Making a Scrollable Frame

Additionally, for users with smaller screens, displaying too much content in the main window can be an issue because not everything will fit. To solve this, we’ve created a frame that fills the entire window, acting as a container for all the GUI elements. The first component we add is the canvas, which will hold the content like images, buttons, and more. The main reason we created this frame is to add a vertical scrollbar, allowing users to scroll when there’s too much content to fit on the screen.

# Create a scrollable frame
main_frame = tk.Frame(root)
main_frame.pack(fill=tk.BOTH, expand=1)

canvas = tk.Canvas(main_frame)
canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=1)

scrollbar = tk.Scrollbar(main_frame, orient=tk.VERTICAL, command=canvas.yview)
scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

canvas.configure(yscrollcommand=scrollbar.set)
canvas.bind('<Configure>', lambda e: canvas.configure(scrollregion=canvas.bbox("all")))

# Create a frame inside the canvas for the content
content_frame = tk.Frame(canvas)
canvas.create_window((0, 0), window=content_frame, anchor="nw")

Managing Global Variables and Device Settings

Now we’ll introduce the VIPs of our program: the global variables that it relies on:

  • selected_image_path: Holds the directory path to the image selected by the user for upscaling.
  • upscaled_img: Stores the upscaled image once the model has worked its magic.
  • device: Defaults to CPU but will later allow the user to choose GPU if available.
# Global variables to store the selected and upscaled images
selected_image_path = None
upscaled_img = None

# Set default device to CPU
device = "cpu"

Loading the Stable Diffusion Upscaler Model for CPU and GPU

Since we’ve got our variables set up, let’s dive into the heart of the operation—the Stable Diffusion Upscaler model.

  • First off, we load a model with float32 precision to run on the CPU. Sure, GPUs are faster, but not everyone has one at their disposal.
pipe_cpu = StableDiffusionUpscalePipeline.from_pretrained(
   "stabilityai/stable-diffusion-x4-upscaler",
   torch_dtype=torch.float32  # Use float32 for CPU
).to("cpu")
  • If you’re lucky enough to have a GPU and want to speed things up, we’ve got float16 precision for you to get those results even quicker!
pipe_gpu = StableDiffusionUpscalePipeline.from_pretrained(
   "stabilityai/stable-diffusion-x4-upscaler",
   torch_dtype=torch.float16  # Use float16 for GPU
).to("cuda") if torch.cuda.is_available() else pipe_cpu
  • And for those sticking with CPU, there’s no need to fret about memory usage. By enabling attention slicing, we ensure the model runs efficiently without hogging too much memory.
pipe_cpu.enable_attention_slicing(None)  # Optimize for CPU performance

Device Configuration and Image Selection Functions

Letting the User Choose

Here, we’ll give users the option to run the program on either the CPU or GPU using a drop-down menu (which we’ll create later) that saves their choice in the device variable.

# Function to select device (CPU or GPU)
def select_device(device_choice):
   global device
   device = device_choice

Picking an Image

At this point, the fun begins! With the select_image() function, we’re giving users the ability to select an image from their computer using filedialog, which opens a window for their selection. Once they’ve chosen an image, the Pillow library works its magic by loading, resizing, and displaying it.

# Function to open a file dialog and select an image
def select_image():
   global selected_image_path, img_label, upscaled_img_label


   file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg *.jpeg *.png")])


   if file_path:
       selected_image_path = file_path
       original_img = Image.open(selected_image_path)
       original_img = original_img.resize((200, 200))  # Resize the image for display


       imgtk = ImageTk.PhotoImage(original_img)
       img_label.config(image=imgtk)
       img_label.image = imgtk

Upscaling the Selected Image and Displaying the Result

Upscaling the Image

So, what are we doing here? Essentially, we’re resizing the image to 128×128 pixels for the CPU to process it more quickly. On the other hand, we maintain the original size for the GPU since it can handle larger images without slowing down.

# Function to upscale the selected image using Stable Diffusion Upscaler
def upscale_image_with_diffusion(image_path):
   global pipe_cpu, pipe_gpu, device


   # Step 1: Open the original image
   image = Image.open(image_path)


   # Check if CPU or GPU is selected
   if device == "cpu":
       # Resize image for CPU to 128x128 for faster performance
       image = image.resize((128, 128))
       upscaled_image = pipe_cpu(prompt="high-quality photo", image=image, num_inference_steps=25).images[0]
   else:
       # For GPU, use the full-resolution image (512x512 or higher)
       image = image.resize((512, 512))
       upscaled_image = pipe_gpu(prompt="very high-quality, highly detailed, accurate photo",
                                 image=image, num_inference_steps=50).images[0]


   return upscaled_image

Based on the chosen device (CPU or GPU), we then pass the image through the appropriate Stable Diffusion Upscaler model. This process generates a higher-quality version of the image using the text prompt “very high-quality, highly detailed, accurate photo”.

Displaying the Upscaled Image

The display_upscaled_image() function is where the magic happens when it comes to showing off our upscaled image. First, we check if the user has actually selected an image. If they have, we disable the upscaling and saving buttons. This way, we avoid any accidental clicks while the processing is underway.

Next up, we want to keep our app feeling snappy and responsive, so instead of running everything on the main thread, we kick off the upscaling process in a separate thread. This means the user can still interact with the app while we work our magic behind the scenes!

# Function to display the upscaled image (runs in the main thread)
def display_upscaled_image():
   global selected_image_path, upscaled_img


   if selected_image_path:
       # Disable buttons while upscaling
       upscale_button.config(state=tk.DISABLED)
       save_button.config(state=tk.DISABLED)


       # Run the upscaling process in a separate thread
       threading.Thread(target=upscale_in_background).start()

Upscaling the Image in the Background and Saving It

Upscaling in the Background

Let’s dive into the upscale_in_background() function, which runs in a separate thread. In the try block, we call the upscale_image_with_diffusion() function we created earlier. This function opens the image and processes it through the Stable Diffusion model.

upscaled_img = upscale_image_with_diffusion(selected_image_path)

Once the image is upscaled, we resize it again to 200×200 pixels to fit the display area.

upscaled_img = upscaled_img.resize((200, 200))  # Resize the upscaled image for display

With the upscaled and resized image ready, its data is stored in the imgtk variable. We then update the GUI by configuring the label with upscaled_img_label.config() to display the upscaled image.

imgtk = ImageTk.PhotoImage(upscaled_img)
upscaled_img_label.config(image=imgtk)
upscaled_img_label.image = imgtk

Finally, once the upscaling and display process is complete, we re-enable the buttons using config(state=tk.NORMAL).

upscale_button.config(state=tk.NORMAL)
save_button.config(state=tk.NORMAL)

Saving the Upscaled Image

Once we’ve upscaled and displayed our image, the next step is to save it, and that’s exactly what the save_image() function is for:

  • It starts by checking for an upscaled image. In the absence of one, a friendly error message will pop up to inform the user. When everything’s in order, a file dialog opens, allowing the user to choose the desired location for saving the image. Upon a successful save, a confirmation message appears to let them know it was saved successfully. In case of any issues, an error message will take its place.
# Function to save the upscaled image
def save_image():
   global upscaled_img


   if upscaled_img:
       save_path = filedialog.asksaveasfilename(defaultextension=".png",
                                                filetypes=[("PNG files", "*.png"), ("JPEG files", "*.jpg"),
                                                           ("All files", "*.*")])
       if save_path:
           upscaled_img.save(save_path)
           messagebox.showinfo("Image Saved", f"Image has been saved successfully at {save_path}")
   else:
       messagebox.showerror("Error", "No upscaled image to save!")

Setting Up the GUI Elements

With the core functions ready, it’s time to set up the layout for our interface. We start by creating the “Select Image” button, which calls the select_image() function. Next, we add a label and a drop-down menu to allow the user to choose between the GPU and CPU, connecting this to the select_device() function.

After that, we create two labels: one for displaying the selected image and another for the upscaled image. Then, we add the “Upscale Image” button that triggers the display_upscaled_image() function and the “Save Image” button that calls the save_image() function.

# Create the layout: Select Image button, Upscale Image button, Save Image button, and labels for images
select_button = Button(content_frame, text="Select Image", command=select_image)
select_button.pack(pady=10)


# Dropdown to select CPU or GPU
device_choice_label = Label(content_frame, text="Choose Device:")
device_choice_label.pack(pady=5)


device_choice = ttk.Combobox(content_frame, values=["cpu", "gpu"], state="readonly")
device_choice.set("cpu")  # Set default value to CPU
device_choice.pack(pady=5)
device_choice.bind("<<ComboboxSelected>>", lambda event: select_device(device_choice.get()))


img_label = Label(content_frame)  # To display the selected image
img_label.pack(pady=10)


upscale_button = Button(content_frame, text="Upscale Image", command=display_upscaled_image)
upscale_button.pack(pady=10)


upscaled_img_label = Label(content_frame)  # To display the upscaled image
upscaled_img_label.pack(pady=10)


save_button = Button(content_frame, text="Save Image", command=save_image)  # Button to save the upscaled image
save_button.pack(pady=10)

Lastly, to keep the main window responsive and running until the user decides to exit, we use the mainloop() method.

# Start the Tkinter event loop
root.mainloop()

Example

This is a low resolution image I’m going to upscale using this code:

This is the result (upscaled image):

Full Code

View on Github
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from tkinter import Label, Button
from PIL import Image, ImageTk
from diffusers import StableDiffusionUpscalePipeline
import torch
import threading


# Initialize the main window
root = tk.Tk()
root.title("Image Upscaler - The Pycodes")
root.geometry("600x400")


# Create a scrollable frame
main_frame = tk.Frame(root)
main_frame.pack(fill=tk.BOTH, expand=1)


canvas = tk.Canvas(main_frame)
canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=1)


scrollbar = tk.Scrollbar(main_frame, orient=tk.VERTICAL, command=canvas.yview)
scrollbar.pack(side=tk.RIGHT, fill=tk.Y)


canvas.configure(yscrollcommand=scrollbar.set)
canvas.bind('<Configure>', lambda e: canvas.configure(scrollregion=canvas.bbox("all")))


# Create a frame inside the canvas for the content
content_frame = tk.Frame(canvas)
canvas.create_window((0, 0), window=content_frame, anchor="nw")


# Global variables to store the selected and upscaled images
selected_image_path = None
upscaled_img = None


# Set default device to CPU
device = "cpu"


# Load the Stable Diffusion Upscaler model for both CPU and GPU
pipe_cpu = StableDiffusionUpscalePipeline.from_pretrained(
   "stabilityai/stable-diffusion-x4-upscaler",
   torch_dtype=torch.float32  # Use float32 for CPU
).to("cpu")


pipe_gpu = StableDiffusionUpscalePipeline.from_pretrained(
   "stabilityai/stable-diffusion-x4-upscaler",
   torch_dtype=torch.float16  # Use float16 for GPU
).to("cuda") if torch.cuda.is_available() else pipe_cpu


pipe_cpu.enable_attention_slicing(None)  # Optimize for CPU performance


# Function to select device (CPU or GPU)
def select_device(device_choice):
   global device
   device = device_choice


# Function to open a file dialog and select an image
def select_image():
   global selected_image_path, img_label, upscaled_img_label


   file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg *.jpeg *.png")])


   if file_path:
       selected_image_path = file_path
       original_img = Image.open(selected_image_path)
       original_img = original_img.resize((200, 200))  # Resize the image for display


       imgtk = ImageTk.PhotoImage(original_img)
       img_label.config(image=imgtk)
       img_label.image = imgtk




# Function to upscale the selected image using Stable Diffusion Upscaler
def upscale_image_with_diffusion(image_path):
   global pipe_cpu, pipe_gpu, device


   # Step 1: Open the original image
   image = Image.open(image_path)


   # Check if CPU or GPU is selected
   if device == "cpu":
       # Resize image for CPU to 128x128 for faster performance
       image = image.resize((128, 128))
       upscaled_image = pipe_cpu(prompt="high-quality photo", image=image, num_inference_steps=25).images[0]
   else:
       # For GPU, use the full-resolution image (512x512 or higher)
       image = image.resize((512, 512))
       upscaled_image = pipe_gpu(prompt="very high-quality, highly detailed, accurate photo",
                                 image=image, num_inference_steps=50).images[0]


   return upscaled_image




# Function to display the upscaled image (runs in the main thread)
def display_upscaled_image():
   global selected_image_path, upscaled_img


   if selected_image_path:
       # Disable buttons while upscaling
       upscale_button.config(state=tk.DISABLED)
       save_button.config(state=tk.DISABLED)


       # Run the upscaling process in a separate thread
       threading.Thread(target=upscale_in_background).start()


# Function to upscale the image in a separate thread
def upscale_in_background():
   global selected_image_path, upscaled_img


   try:
       # Use the Stable Diffusion model to upscale the image
       upscaled_img = upscale_image_with_diffusion(selected_image_path)
       upscaled_img = upscaled_img.resize((200, 200))  # Resize the upscaled image for display


       # Update the UI with the upscaled image (back to the main thread)
       imgtk = ImageTk.PhotoImage(upscaled_img)
       upscaled_img_label.config(image=imgtk)
       upscaled_img_label.image = imgtk
   finally:
       # Re-enable buttons after upscaling is done
       upscale_button.config(state=tk.NORMAL)
       save_button.config(state=tk.NORMAL)




# Function to save the upscaled image
def save_image():
   global upscaled_img


   if upscaled_img:
       save_path = filedialog.asksaveasfilename(defaultextension=".png",
                                                filetypes=[("PNG files", "*.png"), ("JPEG files", "*.jpg"),
                                                           ("All files", "*.*")])
       if save_path:
           upscaled_img.save(save_path)
           messagebox.showinfo("Image Saved", f"Image has been saved successfully at {save_path}")
   else:
       messagebox.showerror("Error", "No upscaled image to save!")




# Create the layout: Select Image button, Upscale Image button, Save Image button, and labels for images
select_button = Button(content_frame, text="Select Image", command=select_image)
select_button.pack(pady=10)


# Dropdown to select CPU or GPU
device_choice_label = Label(content_frame, text="Choose Device:")
device_choice_label.pack(pady=5)


device_choice = ttk.Combobox(content_frame, values=["cpu", "gpu"], state="readonly")
device_choice.set("cpu")  # Set default value to CPU
device_choice.pack(pady=5)
device_choice.bind("<<ComboboxSelected>>", lambda event: select_device(device_choice.get()))


img_label = Label(content_frame)  # To display the selected image
img_label.pack(pady=10)


upscale_button = Button(content_frame, text="Upscale Image", command=display_upscaled_image)
upscale_button.pack(pady=10)


upscaled_img_label = Label(content_frame)  # To display the upscaled image
upscaled_img_label.pack(pady=10)


save_button = Button(content_frame, text="Save Image", command=save_image)  # Button to save the upscaled image
save_button.pack(pady=10)


# Start the Tkinter event loop
root.mainloop()

Happy Coding!

Leave a Comment

Your email address will not be published. Required fields are marked *

Scroll to Top
×