diff --git a/src/constraints/consumer.jl b/src/constraints/consumer.jl
index da186088..ba5c1e76 100644
--- a/src/constraints/consumer.jl
+++ b/src/constraints/consumer.jl
@@ -12,29 +12,26 @@ add_consumer_constraints!(model,
 Adds the consumer asset constraints to the model.
 """
 
-function add_consumer_constraints!(model, constraints, graph, sets)
-    Ac = sets[:Ac]
-    incoming_flow_highest_in_out_resolution = constraints[:highest_in_out].expressions[:incoming]
-    outgoing_flow_highest_in_out_resolution = constraints[:highest_in_out].expressions[:outgoing]
+function add_consumer_constraints!(model, constraints)
+    cons = constraints[:highest_in_out]
+    incoming_flow_highest_in_out_resolution = cons.expressions[:incoming]
+    outgoing_flow_highest_in_out_resolution = cons.expressions[:outgoing]
 
     # - Balance constraint (using the lowest temporal resolution)
-    df = filter(:asset => ∈(Ac), constraints[:highest_in_out].indices; view = true)
     model[:consumer_balance] = [
-        @constraint(
-            model,
-            incoming_flow_highest_in_out_resolution[row.index] -
-            outgoing_flow_highest_in_out_resolution[row.index] -
-            profile_aggregation(
-                Statistics.mean,
-                graph[row.asset].rep_periods_profiles,
-                row.year,
-                row.year,
-                ("demand", row.rep_period),
-                row.time_block_start:row.time_block_end,
-                1.0,
-            ) * graph[row.asset].peak_demand[row.year] in
-            graph[row.asset].consumer_balance_sense,
-            base_name = "consumer_balance[$(row.asset),$(row.year),$(row.rep_period),$(row.time_block_start):$(row.time_block_end)]"
-        ) for row in eachrow(df)
+        begin
+            consumer_balance_sense = if ismissing(row.consumer_balance_sense)
+                MathOptInterface.EqualTo(0.0)
+            else
+                MathOptInterface.GreaterThan(0.0)
+            end
+            @constraint(
+                model,
+                incoming_flow_highest_in_out_resolution[row.index] -
+                outgoing_flow_highest_in_out_resolution[row.index] -
+                row.demand_agg * row.peak_demand in consumer_balance_sense,
+                base_name = "consumer_balance[$(row.asset),$(row.year),$(row.rep_period),$(row.time_block_start):$(row.time_block_end)]"
+            )
+        end for row in eachrow(cons.indices) if row.type == "consumer"
     ]
 end
diff --git a/src/constraints/create.jl b/src/constraints/create.jl
index 56c15a28..c31bd7bd 100644
--- a/src/constraints/create.jl
+++ b/src/constraints/create.jl
@@ -32,19 +32,50 @@ function compute_constraints_indices(connection)
         DuckDB.query(
             connection,
             "CREATE OR REPLACE TEMP SEQUENCE id START 1;
+            CREATE OR REPLACE TABLE cons_highest_in_out AS
             SELECT
                 nextval('id') AS index,
-                t_high.*
+                t_high.*,
+                asset.type,
+                asset.consumer_balance_sense,
+                asset_milestone.peak_demand,
+                COALESCE(mean(profile.value), 1) AS demand_agg,
             FROM t_highest_all_flows AS t_high
             LEFT JOIN asset
                 ON t_high.asset = asset.asset
+            LEFT JOIN asset_milestone
+                ON t_high.asset = asset_milestone.asset
+                AND t_high.year = asset_milestone.milestone_year
             LEFT JOIN asset_both
                 ON t_high.asset = asset_both.asset
                 AND t_high.year = asset_both.milestone_year
                 AND t_high.year = asset_both.commission_year
+            LEFT JOIN (
+                SELECT
+                    assets_profiles.asset,
+                    profile.profile_name,
+                    profile.rep_period,
+                    profile.year,
+                    profile.timestep,
+                    profile.value,
+                FROM assets_profiles
+                LEFT JOIN profiles_rep_periods AS profile
+                    ON profile.profile_name = assets_profiles.profile_name
+                    AND profile.year = assets_profiles.commission_year -- TODO: Is this correct?
+                    AND profile_type = 'demand'
+            ) AS profile
+                ON profile.asset = t_high.asset
+                AND profile.year = t_high.year
+                AND profile.rep_period = t_high.rep_period
+                AND t_high.time_block_start <= profile.timestep
+                AND profile.timestep <= t_high.time_block_end
             WHERE
                 asset_both.active = true
-                AND asset.type in ('hub', 'consumer')",
+                AND asset.type in ('hub', 'consumer')
+            GROUP BY t_high.*;
+
+            FROM cons_highest_in_out
+            ",
         ) |> DataFrame,
     )
 
diff --git a/src/create-model.jl b/src/create-model.jl
index f92c46e2..5641a715 100644
--- a/src/create-model.jl
+++ b/src/create-model.jl
@@ -109,12 +109,7 @@ function create_model(
 
     @timeit to "add_energy_constraints!" add_energy_constraints!(model, constraints, graph)
 
-    @timeit to "add_consumer_constraints!" add_consumer_constraints!(
-        model,
-        constraints,
-        graph,
-        sets,
-    )
+    @timeit to "add_consumer_constraints!" add_consumer_constraints!(model, constraints)
 
     @timeit to "add_storage_constraints!" add_storage_constraints!(
         model,
diff --git a/test/test-case-studies.jl b/test/test-case-studies.jl
index fc8cdb1b..ef2f0e67 100644
--- a/test/test-case-studies.jl
+++ b/test/test-case-studies.jl
@@ -85,8 +85,16 @@ end
     dir = joinpath(INPUT_FOLDER, "Tiny")
     connection = DBInterface.connect(DuckDB.DB)
     _read_csv_folder(connection, dir)
+    DuckDB.execute(
+        connection,
+        "UPDATE asset_milestone
+            SET peak_demand = -1
+            WHERE
+                asset = 'demand'
+                AND milestone_year = 2030
+        ",
+    )
     energy_problem = EnergyProblem(connection)
-    energy_problem.graph["demand"].peak_demand[2030] = -1 # make it infeasible
     create_model!(energy_problem)
     @test_logs (:warn, "Model status different from optimal") solve_model!(energy_problem)
     @test energy_problem.termination_status == JuMP.INFEASIBLE