Python AttributeError: 'NoneType' object has no attribute 'load_state_dict' when using RealESRGAN model

Question

I'm in the process of building a backend with Python to upscale images using the RealESRGAN model from this git repository: https://github.com/xinntao/Real-ESRGAN.

This is a simplified app.py file in which I'm running tests:

import os
from flask import Flask, request, jsonify
from flask_cors import CORS
from realesrgan import RealESRGANer
from PIL import Image
import numpy as np
import io

app = Flask(__name__)
CORS(app)

# Initialize the RealESRGAN model
model_path = "./RealESRGAN_x4plus_anime_6B.pth"  # Path to the model that can be dowloaded here: https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth

model = RealESRGANer(scale=4, model_path=model_path)  # Initialize the model with the specified scale

@app.route('/upscale', methods=['POST'])
def upscale_image():
    if 'image' not in request.files:
        return jsonify({"error": "No image provided"}), 400
    
    file = request.files['image']

    try:
        # Read the image file
        img = Image.open(file.stream).convert('RGB')

        # Convert image to numpy array and upscale
        img_array = np.array(img)
        upscaled_image = model.predict(img_array)

        # Convert upscaled image back to PIL format
        upscaled_image = Image.fromarray(upscaled_image)

        # Save the upscaled image to a byte stream
        img_byte_arr = io.BytesIO()
        upscaled_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)

        return jsonify({"message": "Image upscaled successfully", "upscaled_image": img_byte_arr.getvalue().decode('latin1')})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True)

However, I'm experiencing the following error when trying to run the app:

model.load_state_dict(loadnet[keyname], strict=True)
    ^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'load_state_dict'

Any suggestions on how to solve this problem?

Here are a few things that I've tried to do while debugging:

Loading of Model: I have tried "manually" loading the model with the following code:

# Detect the device (GPU or CPU)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Initialize RealESRGANer
        self.model = RealESRGANer(
            scale=4,  # Upscale factor
            model_path=model_path,  # Model path for pretrained weights
            tile=0,  # No tiling
            tile_pad=10,  # Padding if tiling is used
            pre_pad=0,  # Pre-padding
            half=False,  # Avoid half precision issues
            device=device  # Use detected device (CUDA or CPU)
        )

        # Load the weights into the model architecture
        state_dict = torch.load(model_path, map_location=device)

        # Ensure 'params_ema' is present in the state_dict
        if "params_ema" in state_dict:
            self.model.model.load_state_dict(state_dict["params_ema"], strict=True)
        else:
            raise ValueError("The state_dict does not contain 'params_ema'. Please check the model file.")

        # Set the model to evaluation mode (important for inference)
        self.model.model.eval()

File Path: I have also included the following code to ensure that my model path is correct:

print(os.path.exists("./RealESRGAN_x4plus_anime_6B.pth"))

Alternative Model: I also have tested the alternative model that can be found here: https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth

mhenning · Accepted Answer · 2024-10-22 14:47:10Z

The problem is that you try to initiate a helper method, which needs the "real" model as model parameter. Default to that model parameter is None, which explains the error message. You can see it in the init of RealESRGANer. You can also see how the actual model is loaded in the inference script in the lines below.
You could copy the model loading line, if the weights you want to use are fixed:

model_path = "./RealESRGAN_x4plus_anime_6B.pth"
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)  # for `RealESRGAN_x4plus_anime_6B`
model = RealESRGANer(scale=4, model=model, model_path=model_path)

Note that in the inference script, if the --face_enhance parameter is used, another model is loaded (code) and the RealESRGANer model is used as a parameter (there it is called upsampler).

Also please note that I could not test my solution, it is just based on digging through the source code.

Dieter Roelofse · Accepted Answer · 2024-10-23 09:39:07Z

I would suggest using the inference_realesrgan.py file found in the official Real-ESRGAN repo to load the model and do the upscaling.

Download the inference_realesrgan.py file and put it in your codebase (same directory as your app.py file), then for example you can create a wrapper function (let's call it inference_wrapper.py) to Call the script like so:

import subprocess
import os

def upscale_image(
    input_image_path,
    output_image_path,
    model_name,
    outscale,
):

    # Build the command to run the inference script
    command = [
        "python",
        "inference_realesrgan.py",  # Path to your script
        "--input",
        input_image_path,
        "--output",
        os.path.dirname(output_image_path),
        "--model_name",
        model_name,
        "--outscale",
        str(outscale),
        "--suffix",
        "upscaled",
    ]

    subprocess.run(command, check=True)

    # Return the path to the upscaled image
    img_name, ext = os.path.splitext(os.path.basename(input_image_path))
    upscaled_image = os.path.join(
        os.path.dirname(output_image_path), f"{img_name}_upscaled{ext}"
    )
    return upscaled_image

then in your app.py file, call the upscale_image function with the appropriate arguments like so:

from inference_wrapper import (
    upscale_image,
)

input_image = 'path/to/input/image.jpg'
output_image = 'path/to/output/folder/image.jpg'
model_name = "RealESRGAN_x4plus_anime_6B" # or whichever model you want to use
outscale = 4

# Call the function to upscale the image
upscaled_image = upscale_image(input_image, output_image, model_name, outscale)

# The image has been upscaled and saved to `upscaled_image`
print(f"Upscaled image saved at: {upscaled_image}")

This should do the trick.

Also: I had a look at the inference_realesrgan.py file and ideally you would want to refactor it a bit to make it a bit more modular, like separating the functionality to create independent functions to:

load the model
download the model (weights)
upscale the input image
enable/disable face enhancer
set the gpu_id
set the half-precision to true/false to use fp16/fp32

Hope you come right.

Collectives™ on Stack Overflow

Python AttributeError: 'NoneType' object has no attribute 'load_state_dict' when using RealESRGAN model

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related