Whenever I try to load the model it gives me the error of only input tensors can be passed as positional arguments.
I dont wanna change the concept of joining layers but i want to load the model and also to have it accept images.

Also need help with the data generators on how they work and how can i implement them into my code or should I manually load the images as an array and pass it into the model.

Here is the complete code :

# Defining U-net architecture with encoder and decoder blocks

def encoder_block(x, filters, kernel_size = (3,3), padding = "valid"):

    x = keras.layers.Conv2D(filters = filters, kernel_size = kernel_size, padding = padding)
    x = keras.layers.Activation(keras.activations.relu)(x)
    x = keras.layers.Dropout(0.5)(x)

    x = keras.layers.Conv2D(filters = filters, kernel_size = kernel_size, padding = padding)(x)
    x = keras.layers.Activation(keras.activations.relu)(x)

    skip_connection = x

    x = tf.keras.layers.MaxPool2D(pool_size = (2, 2), 
                                  strides = 2)(x)
    
    return (x,skip_connection) 

def decoder_block(x, skip_connection, filters, kernel_size = (3,3), padding = "valid"):

    x = keras.layers.Conv2DTranspose(filters, kernel_size = (2,2),
                                     strides=2, padding = padding)
    
    x = tf.keras.layers.Concatenate()([x, skip_connection])
    
    x = keras.layers.Conv2D(filters = filters, kernel_size = kernel_size, padding = padding)
    x = keras.layers.Activation(keras.activations.relu)(x)
    x = keras.layers.Dropout(0.5)(x)

    x = keras.layers.Conv2D(filters = filters, kernel_size = kernel_size, padding = padding)(x)
    x = keras.layers.Activation(keras.activations.relu)(x)

    return x

# Structuring the U-net model
# Encoder block : e, skip connection block: skip_conn, decoder block : d

IMAGE_SIZE = (512, 512)
input = keras.layers.Input(shape=IMAGE_SIZE)

# Encoder Path
e1, skip_conn1 = encoder_block(input, filters=64)
e2, skip_conn2 = encoder_block(e1, filters=128)
e3, skip_conn3 = encoder_block(e2, filters=256)
e4, skip_conn4 = encoder_block(e3, filters=512)

# Bottle-neck 
b1 = keras.layers.Conv2D(filters=1024, kernel_size=(3,3), padding="valid")(e4)
b1 = keras.layers.Activation(keras.activations.relu)(b1)
b1 = keras.layers.Dropout(0.5)(b1)
b1 = keras.layers.Conv2D(filters=1024, kernel_size=(3,3), padding="valid")(b1)
b1 = keras.layers.Activation(keras.activations.relu)(b1)

# Decoder Path
d1 = decoder_block(b1, skip_conn4, filters=512)
d2 = decoder_block(d1, skip_conn3, filters=256)
d3 = decoder_block(d2, skip_conn2, filters=128)
d4 = decoder_block(d3, skip_conn1, filters=64)

# Output Layer
Segmentation_output = tf.keras.layers.Conv2D(1, 1, padding = 'valid', activation = 'sigmoid', name = "Segmenmtation")(d4)

classification_output = keras.layers.Flatten()(d4)
classification_output = keras.layers.Dense(7, activation="softmax", name = "Classification")(classification_output)

model = keras.models.Model(inputs=input, outputs=[Segmentation_output, classification_output])