Request: When upgrading images in batch, you can set a fixed resolution that allows dynamic scaling regardless of the size of the input image.

[RexLeeGrey]

When upgrading images in batch, you can set a fixed resolution that allows dynamic scaling regardless of the size of the input image.

[JarJarBeatyourattitude]

Try adding something like this:

# Function to adjust image size and upscale factor
def resize_image_for_upscale(original_image, target_pixels):
    """
    Adjusts the image size and calculates the upscale factor so that
    after upscaling, the image is as close as possible to the target pixel count.

    Args:
        original_image (Image): The original PIL Image.
        target_pixels (int): The target pixel count after upscaling.

    Returns:
        Image: Adjusted image.
        int: Calculated upscale factor.
    """
    original_width, original_height = original_image.size
    print(f"Original image size: Width={original_width}, Height={original_height}")

    current_pixels = original_width * original_height
    print(f"Original pixel count: {current_pixels}")

    upscale_factor = 1
    # Calculate upscale factor that just exceeds the target pixel count
    while current_pixels * (upscale_factor ** 2) < target_pixels:
        upscale_factor += 1

    print(f"Calculated minimum upscale factor: {upscale_factor} to exceed target pixels: {target_pixels}")

    # Calculate new dimensions to closely match the target pixel count with the determined upscale factor
    adjusted_pixel_target = target_pixels / (upscale_factor ** 2)
    print(f"Adjusted pixel target before upscaling: {adjusted_pixel_target}")

    ratio = original_width / original_height
    adjusted_width = int(math.sqrt(adjusted_pixel_target * ratio))
    adjusted_height = int(adjusted_pixel_target / math.sqrt(adjusted_pixel_target * ratio))

    print(f"Adjusted image size to closely match the target pixel count: Width={adjusted_width}, Height={adjusted_height}")

    adjusted_image = original_image.resize((adjusted_width, adjusted_height), Image.ANTIALIAS)
    adjusted_pixel_count = adjusted_width * adjusted_height
    print(f"Adjusted pixel count before upscaling: {adjusted_pixel_count}")
    print(f"Expected pixel count after applying upscale factor: {adjusted_pixel_count * (upscale_factor ** 2)}")

    return adjusted_image, upscale_factor


# load SUPIR
model = create_SUPIR_model('options/SUPIR_v0.yaml', SUPIR_sign=args.SUPIR_sign).to(SUPIR_device)
model = model.half() #TODO Added manually, this halfs the size of the model
if args.use_tile_vae:
    model.init_tile_vae(encoder_tile_size=1800, decoder_tile_size=64)

model.ae_dtype = convert_dtype(args.ae_dtype)
model.model.dtype = convert_dtype(args.diff_dtype)
# load LLaVA
if use_llava:
    llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device)
else:
    llava_agent = None

os.makedirs(args.save_dir, exist_ok=True)
for img_pth in os.listdir(args.img_dir):
    img_name = os.path.splitext(img_pth)[0]

    LQ_img = Image.open(os.path.join(args.img_dir, img_pth))
    # Resize the image before processing if needed
    target_resolution =  3600 * 2000  # Assuming 4K target pixel count. This is as is to ensure 2x2 tiling even at 16:9 or 9:16 ratio, when one side is 3600 pixels. Max without tiling is 1800 * 1800
    LQ_img, upscale_factor = resize_image_for_upscale(LQ_img, target_resolution)

[RexLeeGrey]

Can you tell me which file I need to add these to?

[JarJarBeatyourattitude]

This would go in the test.py if that's what you're using to make the images.